IndexingUtils.cpp Source File (original) (raw)

15#include "llvm/ADT/STLExtras.h"

16#include

17#include

19using namespace mlir;

21template

23 ExprType unit) {

24 if (sizes.empty())

25 return {};

27 for (int64_t r = static_cast<int64_t>(strides.size()) - 2; r >= 0; --r)

28 strides[r] = strides[r + 1] * sizes[r + 1];

29 return strides;

30}

32template

36 if (v1.empty() && v2.empty())

37 return {};

39 for (auto it : llvm::zip_equal(v1, v2))

40 result.push_back(std::get<0>(it) * std::get<1>(it));

42}

44template

47 assert(offsets.size() == basis.size());

48 ExprType linearIndex = zero;

49 for (unsigned idx = 0, e = basis.size(); idx < e; ++idx)

50 linearIndex = linearIndex + offsets[idx] * basis[idx];

51 return linearIndex;

52}

54template <typename ExprType, typename DivOpTy>

57 DivOpTy divOp) {

58 int64_t rank = strides.size();

60 for (int64_t r = 0; r < rank; ++r) {

61 offsets[r] = divOp(linearIndex, strides[r]);

62 linearIndex = linearIndex % strides[r];

63 }

64 return offsets;

65}

72 assert((sizes.empty() ||

73 llvm::all_of(sizes.drop_front(), [](int64_t s) { return s >= 0; })) &&

74 "sizes must be nonnegative");

76 return ::computeSuffixProductImpl(sizes, unit);

77}

85 assert(llvm::all_of(basis, [](int64_t s) { return s > 0; }) &&

86 "basis must be nonnegative");

87 return llvm::product_of(basis);

88}

91 assert(llvm::all_of(basis, [](int64_t s) { return s > 0; }) &&

92 "basis must be nonnegative");

95}

99 assert(llvm::all_of(strides, [](int64_t s) { return s > 0; }) &&

100 "strides must be nonnegative");

103}

104

105std::optional<SmallVector<int64_t>>

107 if (shape.size() < subShape.size())

108 return std::nullopt;

109 assert(llvm::all_of(shape, [](int64_t s) { return s > 0; }) &&

110 "shape must be nonnegative");

111 assert(llvm::all_of(subShape, [](int64_t s) { return s > 0; }) &&

112 "subShape must be nonnegative");

113

114

115 std::vector<int64_t> result;

117 for (auto [size, subSize] :

118 llvm::zip(llvm::reverse(shape), llvm::reverse(subShape))) {

119

120 if (size % subSize != 0)

121 return std::nullopt;

122 result.push_back(size / subSize);

123 }

124

125

126 int commonSize = subShape.size();

127 std::copy(shape.rbegin() + commonSize, shape.rend(),

128 std::back_inserter(result));

129

131}

132

133

134

135

136

138 if (sizes.empty())

139 return {};

141 return ::computeSuffixProductImpl(sizes, unit);

142}

143

148

152

156

162

168

172 linearIndex, strides,

174}

175

181

182

183

184

185

188 assert(llvm::all_of(permutation, [](int64_t s) { return s >= 0; }) &&

189 "permutation must be non-negative");

191 for (const auto &pos : llvm::enumerate(permutation)) {

192 inversion[pos.value()] = pos.index();

193 }

194 return inversion;

195}

196

198 for (auto i : llvm::seq<int64_t>(0, permutation.size()))

199 if (permutation[i] != i)

200 return false;

201 return true;

202}

203

205 llvm::SmallDenseSet<int64_t, 4> seenVals;

206 for (auto val : interchange) {

207 if (val < 0 || static_cast<uint64_t>(val) >= interchange.size())

208 return false;

209 if (seenVals.count(val))

210 return false;

211 seenVals.insert(val);

212 }

213 return seenVals.size() == interchange.size();

214}

215

221 for (auto [pos, desiredPos] : llvm::zip_equal(positions, desiredPositions)) {

222 res[desiredPos] = pos;

223 seen.insert(pos);

224 }

226 for (int64_t &entry : res) {

227 if (entry != -1)

228 continue;

229 while (seen.contains(nextPos))

230 ++nextPos;

231 entry = nextPos;

232 ++nextPos;

233 }

234 return res;

235}

236

239 assert(inputPerm.size() >= dropPositions.size() &&

240 "expect inputPerm size large than position to drop");

242 unsigned permSize = inputPerm.size();

243 for (unsigned inputIndex = 0; inputIndex < permSize; ++inputIndex) {

244 int64_t targetIndex = inputPerm[inputIndex];

245 bool shouldDrop = false;

246 unsigned dropSize = dropPositions.size();

247 for (unsigned dropIndex = 0; dropIndex < dropSize; dropIndex++) {

248 if (dropPositions[dropIndex] == inputPerm[inputIndex]) {

249 shouldDrop = true;

250 break;

251 }

252 if (dropPositions[dropIndex] < inputPerm[inputIndex]) {

253 targetIndex--;

254 }

255 }

256 if (!shouldDrop) {

257 res.push_back(targetIndex);

258 }

259 }

260 return res;

261}

262

265 unsigned dropBack) {

266 assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds");

267 auto range = arrayAttr.getAsRange();

269 res.reserve(arrayAttr.size() - dropFront - dropBack);

270 for (auto it = range.begin() + dropFront, eit = range.end() - dropBack;

271 it != eit; ++it)

272 res.push_back((*it).getValue().getSExtValue());

273 return res;

274}

275

276

278 assert(val && "Invalid value");

279 if (auto attr = dyn_cast(val)) {

280 return attr.getContext();

281 }

282 return cast(val).getContext();

283}

284

285std::pair<AffineExpr, SmallVector>

289 assert(strides.size() == indices.size());

290 auto sourceRank = static_cast<unsigned>(strides.size());

291

292

295

298 values[0] = sourceOffset;

299

300 for (unsigned i = 0; i < sourceRank; ++i) {

301

303

304

305 unsigned baseIdxForDim = 1 + 2 * i;

306 unsigned subOffsetForDim = baseIdxForDim;

307 unsigned origStrideForDim = baseIdxForDim + 1;

308 expr = expr + symbols[subOffsetForDim] * symbols[origStrideForDim];

309 values[subOffsetForDim] = indices[i];

310 values[origStrideForDim] = origStride;

311 }

312

313 return {expr, values};

314}

315

316std::pair<AffineExpr, SmallVector>

323

324

325

326

327

328

330 unsigned paddedSize) {

331 assert(tileShape.size() <= paddedSize &&

332 "expected tileShape to <= paddedSize");

333 if (tileShape.size() == paddedSize)

334 return to_vector(tileShape);

336 llvm::append_range(result, tileShape);

338}

339

345 sliceStrides(shape.size()) {

346

347 std::optional<SmallVector<int64_t>> shapeRatio =

349 assert(shapeRatio && shapeRatio->size() == shape.size() &&

350 "target shape does not evenly divide the original shape");

352 "expected loop order to be a permutation of rank equal to outer "

353 "shape");

354

358}

359

361 int64_t linearIndex) const {

363 delinearize(linearIndex, sliceStrides), inverseLoopOrder);

365}

366

372 delinearize(linearIndex, sliceStrides), inverseLoopOrder);

375}

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

static SmallVector< ExprType > delinearizeImpl(ExprType linearIndex, ArrayRef< ExprType > strides, DivOpTy divOp)

Definition IndexingUtils.cpp:55

static SmallVector< ExprType > computeElementwiseMulImpl(ArrayRef< ExprType > v1, ArrayRef< ExprType > v2)

Definition IndexingUtils.cpp:33

static SmallVector< int64_t > padTileShapeToSize(ArrayRef< int64_t > tileShape, unsigned paddedSize)

Apply left-padding by 1 to the tile shape if required.

Definition IndexingUtils.cpp:329

static ExprType linearizeImpl(ArrayRef< ExprType > offsets, ArrayRef< ExprType > basis, ExprType zero)

Definition IndexingUtils.cpp:45

static SmallVector< ExprType > computeSuffixProductImpl(ArrayRef< ExprType > sizes, ExprType unit)

Definition IndexingUtils.cpp:22

Base type for affine expression.

AffineExpr floorDiv(uint64_t v) const

MLIRContext * getContext() const

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

This class represents a single result from folding an operation.

MLIRContext * getContext() const

SmallVector< int64_t > getStaticTileOffsets(int64_t linearIndex) const

Definition IndexingUtils.cpp:360

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

Definition IndexingUtils.cpp:340

SmallVector< AffineExpr > getDynamicTileOffsets(AffineExpr linearIndex) const

Definition IndexingUtils.cpp:368

Include the generated interface declarations.

OpFoldResult getAsIndexOpFoldResult(MLIRContext *ctx, int64_t val)

Convert int64_t to integer attributes of index type and return them as OpFoldResult.

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

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

Definition IndexingUtils.cpp:79

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

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.

Definition IndexingUtils.cpp:286

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...

Definition IndexingUtils.cpp:97

llvm::DenseSet< ValueT, ValueInfoT > DenseSet

int64_t computeProduct(ArrayRef< int64_t > basis)

Self-explicit.

Definition IndexingUtils.cpp:84

bool isIdentityPermutation(ArrayRef< int64_t > permutation)

Returns true if permutation is an identity permutation.

Definition IndexingUtils.cpp:197

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...

Definition IndexingUtils.cpp:217

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.

Definition IndexingUtils.cpp:263

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

Given a set of sizes, return the suffix product.

Definition IndexingUtils.cpp:71

int64_t computeMaxLinearIndex(ArrayRef< int64_t > basis)

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

AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)

OpFoldResult getAsOpFoldResult(Value val)

Given a value, try to extract a constant Attribute.

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

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

Definition IndexingUtils.cpp:90

SmallVector< AffineExpr > getAffineConstantExprs(ArrayRef< int64_t > constants, MLIRContext *context)

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.

Definition IndexingUtils.cpp:106

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.

Definition IndexingUtils.cpp:237

bool isPermutationVector(ArrayRef< int64_t > interchange)

Method to check if an interchange vector is a permutation.

Definition IndexingUtils.cpp:204

void bindSymbolsList(MLIRContext *ctx, MutableArrayRef< AffineExprTy > exprs)

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

Helper method to apply to inverse a permutation.

Definition IndexingUtils.cpp:187