IndexingUtils.h Source File (original) (raw)

14 #ifndef MLIR_DIALECT_UTILS_INDEXINGUTILS_H

15 #define MLIR_DIALECT_UTILS_INDEXINGUTILS_H

19 #include "llvm/ADT/ArrayRef.h"

20 #include "llvm/ADT/SmallVector.h"

21 #include "llvm/ADT/iterator.h"

22 #include

23 #include

25 namespace mlir {

26 class ArrayAttr;

49 }

55 ArrayRef<int64_t> v2);

58 int64_t computeSum(ArrayRef<int64_t> basis);

71 }

76 int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis);

85 SmallVector<int64_t> delinearize(int64_t linearIndex,

86 ArrayRef<int64_t> strides);

100

101

102

103 std::optional<SmallVector<int64_t>>

104 computeShapeRatio(ArrayRef<int64_t> shape, ArrayRef<int64_t> subShape);

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

131 }

132

133

134

135

136

137

138

139

141 ArrayRef v2);

142

143

144 AffineExpr computeSum(MLIRContext *ctx, ArrayRef basis);

145

146

147 AffineExpr computeProduct(MLIRContext *ctx, ArrayRef basis);

148

149

150

151

152

153

154

155

156

157

158

162 }

163

164

165

166

167

168

169

170

171

172

173

174 AffineExpr linearize(MLIRContext *ctx, ArrayRef offsets,

175 ArrayRef basis);

176 AffineExpr linearize(MLIRContext *ctx, ArrayRef offsets,

177 ArrayRef<int64_t> basis);

178

179

180

181

182

183

184

185

186

187

188

189

190

191 SmallVector delinearize(AffineExpr linearIndex,

192 ArrayRef strides);

193 SmallVector delinearize(AffineExpr linearIndex,

194 ArrayRef<int64_t> strides);

195

196

197

198

199

200 template

203 assert(input.size() == permutation.size() &&

204 "expected input rank to equal permutation rank");

205 assert(

206 llvm::all_of(permutation, [&](size_t s) { return s < input.size(); }) &&

207 "permutation must be within input bounds");

208 auto permutationRange = llvm::map_range(

209 llvm::seq(0, input.size()),

210 [&](int64_t idx) -> T { return input[permutation[idx]]; });

211 return llvm::to_vector(permutationRange);

212 }

213

214 template

218 }

219

220

221

222

223

224

225 template <typename T, unsigned N>

229 }

230

231

233

234

236

237

239

240

241

242

243

244

245 SmallVector<int64_t>

247 ArrayRef<int64_t> desiredPositions);

248

249

250

251

252

253

254 SmallVector<int64_t> dropDims(ArrayRef<int64_t> inputPerm,

255 ArrayRef<int64_t> dropPositions);

256

257

258

260 unsigned dropBack = 0);

261

262

263

264

265

266

267

268 std::pair<AffineExpr, SmallVector>

269 computeLinearIndex(OpFoldResult sourceOffset, ArrayRef strides,

270 ArrayRef indices);

271 std::pair<AffineExpr, SmallVector>

272 computeLinearIndex(OpFoldResult sourceOffset, ArrayRef<int64_t> strides,

273 ArrayRef indices);

274

275

276

277

278

279 namespace detail {

280

281

283 public:

286

288

290

292

293 template

295 if constexpr (std::is_same_v<T, int64_t>)

297 else

299 }

300

301 size_t getRank() const { return tileShape.size(); }

302

303 private:

304

305

307

308

310

312

313

314 int64_t maxLinearIndex;

315 };

316

317

318 template

320 : public llvm::iterator_facade_base<TileOffsetRangeIterator,

321 std::forward_iterator_tag,

322 SmallVector> {

323 public:

325 : params(params), index(index) {}

326

329 const auto copy = *this;

330 ++*this;

332 }

333

335 return index == other.index;

336 }

338 return index != other.index;

339 }

340

343 }

344 void operator+=(int64_t offset) { incrementIndex(offset); }

345

346 private:

347 void incrementIndex(int64_t offset) { index = index + offset; }

348 const TileOffsetRangeImpl params;

349 int64_t index;

350 };

351 }

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

377 public:

380

383 : params(shape, tileShape, loopOrder), beginValue(params, 0),

384 pastEndValue(params, params.getMaxLinearIndex()) {

385 assert(shape.size() >= tileShape.size());

386 assert(loopOrder.size() == shape.size());

387 }

388

389

391 : params(shape, tileShape,

392 llvm::to_vector(llvm::seq<int64_t>(0, shape.size()))),

393 beginValue(params, 0),

394 pastEndValue(params, params.getMaxLinearIndex()) {

395 assert(shape.size() >= tileShape.size());

396 }

397

400

401

403

404

406

407 private:

411 };

412 }

413

414 #endif

void dropFront(int64_t arr[N], int64_t *res)

static void copy(Location loc, Value dst, Value src, Value size, OpBuilder &builder)

Copies the given number of bytes from src to dst pointers.

Base type for affine expression.

MLIRContext is the top-level object for a collection of MLIR operations.

A range-style iterator that allows for iterating over the offsets of all potential tiles of size tile...

size_t size() const

Returns the total number of tiles that fit in the larger shape.

size_t getRank() const

Returns rank of the iterator's shape.

StaticTileOffsetRange(ArrayRef< int64_t > shape, ArrayRef< int64_t > tileShape, ArrayRef< int64_t > loopOrder)

detail::TileOffsetRangeIterator< int64_t > IteratorTy

detail::TileOffsetRangeImpl ParamsTy

StaticTileOffsetRange(ArrayRef< int64_t > shape, ArrayRef< int64_t > tileShape)

Create the range with identity loop order.

Encapsulates the set of parameters that are used to make tile offset calculations in the TileOffsetRa...

int64_t getMaxLinearIndex() const

SmallVector< int64_t > getStaticTileOffsets(int64_t linearIndex) const

TileOffsetRangeImpl(ArrayRef< int64_t > shape, ArrayRef< int64_t > tileShape, ArrayRef< int64_t > loopOrder)

SmallVector< T > getTileOffsets(T linearIndex) const

SmallVector< AffineExpr > getDynamicTileOffsets(AffineExpr linearIndex) const

The STL-style iterator implementation for StaticTileOffsetRange.

TileOffsetRangeIterator(const TileOffsetRangeImpl &params, ElementType index)

bool operator!=(const TileOffsetRangeIterator &other) const

void operator+=(int64_t offset)

SmallVector< ElementType > operator*() const

TileOffsetRangeIterator operator++(int)

bool operator==(const TileOffsetRangeIterator &other) const

The OpAsmOpInterface, see OpAsmInterface.td for more details.

Include the generated interface declarations.

SmallVector< int64_t > computeElementwiseMul(ArrayRef< int64_t > v1, ArrayRef< int64_t > v2)

Return a vector containing llvm::zip_equal(v1, v2) multiplied elementwise.

std::pair< AffineExpr, SmallVector< OpFoldResult > > computeLinearIndex(OpFoldResult sourceOffset, ArrayRef< OpFoldResult > strides, ArrayRef< OpFoldResult > indices)

Compute linear index from provided strides and indices, assuming strided layout.

SmallVector< int64_t > computeStrides(ArrayRef< int64_t > sizes)

SmallVector< T > applyPermutation(ArrayRef< T > input, ArrayRef< int64_t > permutation)

SmallVector< int64_t > delinearize(int64_t linearIndex, ArrayRef< int64_t > strides)

Given the strides together with a linear index in the dimension space, return the vector-space offset...

int64_t computeProduct(ArrayRef< int64_t > basis)

Self-explicit.

bool isIdentityPermutation(ArrayRef< int64_t > permutation)

Returns true if permutation is an identity permutation.

SmallVector< int64_t > computePermutationVector(int64_t permSize, ArrayRef< int64_t > positions, ArrayRef< int64_t > desiredPositions)

Return a permutation vector of size permSize that would result in moving positions into desiredPositi...

SmallVector< int64_t > getI64SubArray(ArrayAttr arrayAttr, unsigned dropFront=0, unsigned dropBack=0)

Helper to return a subset of arrayAttr as a vector of int64_t.

SmallVector< int64_t > computeSuffixProduct(ArrayRef< int64_t > sizes)

Given a set of sizes, return the suffix product.

int64_t computeMaxLinearIndex(ArrayRef< int64_t > basis)

Return the number of elements of basis (i.e.

int64_t linearize(ArrayRef< int64_t > offsets, ArrayRef< int64_t > basis)

Return the linearized index of 'offsets' w.r.t.

std::optional< SmallVector< int64_t > > computeShapeRatio(ArrayRef< int64_t > shape, ArrayRef< int64_t > subShape)

Return the multi-dimensional integral ratio of subShape to the trailing dimensions of shape.

void applyPermutationToVector(SmallVector< T, N > &inVec, ArrayRef< int64_t > permutation)

Apply the permutation defined by permutation to inVec.

int64_t computeSum(ArrayRef< int64_t > basis)

Self-explicit.

SmallVector< int64_t > dropDims(ArrayRef< int64_t > inputPerm, ArrayRef< int64_t > dropPositions)

Returns a permutation vector that drop the input dims in dropPositions from inputPerm.

bool isPermutationVector(ArrayRef< int64_t > interchange)

Method to check if an interchange vector is a permutation.

SmallVector< int64_t > invertPermutationVector(ArrayRef< int64_t > permutation)

Helper method to apply to inverse a permutation.