22 for (int64_t r = strides.size() - 2; r >= 0; --r)
23 strides[r] = strides[r + 1] * sizes[r + 1];
30 for (
auto it : llvm::zip(v1, v2))
31 result.push_back(std::get<0>(it) * std::get<1>(it));
35 std::optional<SmallVector<int64_t>>
37 if (shape.size() < subShape.size())
39 assert(llvm::all_of(shape, [](int64_t s) {
return s > 0; }) &&
40 "shape must be nonnegative");
41 assert(llvm::all_of(subShape, [](int64_t s) {
return s > 0; }) &&
42 "subShape must be nonnegative");
45 std::vector<int64_t> result;
46 result.reserve(shape.size());
47 for (
auto [size, subSize] :
48 llvm::zip(llvm::reverse(shape), llvm::reverse(subShape))) {
50 if (size % subSize != 0)
52 result.push_back(size / subSize);
56 int commonSize = subShape.size();
57 std::copy(shape.rbegin() + commonSize, shape.rend(),
58 std::back_inserter(result));
64 assert(offsets.size() == basis.size());
65 int64_t linearIndex = 0;
66 for (
unsigned idx = 0, e = basis.size(); idx < e; ++idx)
67 linearIndex += offsets[idx] * basis[idx];
73 int64_t rank = sliceStrides.size();
75 for (int64_t r = 0; r < rank; ++r) {
76 assert(sliceStrides[r] > 0);
77 vectorOffsets[r] = index / sliceStrides[r];
78 index %= sliceStrides[r];
86 return std::accumulate(basis.begin(), basis.end(), 1,
87 std::multiplies<int64_t>());
94 inversion[pos.value()] = pos.index();
100 llvm::SmallDenseSet<int64_t, 4> seenVals;
101 for (
auto val : interchange) {
102 if (seenVals.count(val))
104 seenVals.insert(val);
106 return seenVals.size() == interchange.size();
112 assert(arrayAttr.size() >
dropFront + dropBack &&
"Out of bounds");
113 auto range = arrayAttr.getAsRange<IntegerAttr>();
115 res.reserve(arrayAttr.size() -
dropFront - dropBack);
116 for (
auto it = range.begin() +
dropFront, eit = range.end() - dropBack;
118 res.push_back((*it).getValue().getSExtValue());
125 resultExpr = resultExpr * basis[0];
126 for (
unsigned i = 1; i < basis.size(); i++)
134 int64_t rank = strides.size();
136 vectorOffsets[0] = resultExpr.
floorDiv(strides[0]);
137 resultExpr = resultExpr % strides[0];
138 for (
unsigned i = 1; i < rank; i++) {
139 vectorOffsets[i] = resultExpr.
floorDiv(strides[i]);
140 resultExpr = resultExpr % strides[i];
142 return vectorOffsets;
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.
AffineExpr floorDiv(uint64_t v) const
This class is a general helper class for creating context-global objects like types,...
AffineExpr getAffineDimExpr(unsigned position)
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Include the generated interface declarations.
SmallVector< int64_t > computeElementwiseMul(ArrayRef< int64_t > v1, ArrayRef< int64_t > v2)
Return a vector containing llvm::zip of v1 and v2 multiplied elementwise.
SmallVector< int64_t > computeStrides(ArrayRef< int64_t > sizes)
Given a set of sizes, compute and return the strides (i.e.
SmallVector< mlir::AffineExpr > getDelinearizedAffineExpr(ArrayRef< int64_t > strides, mlir::Builder &b)
Given the strides in the dimension space, returns the affine expressions for vector-space offsets in ...
SmallVector< int64_t > delinearize(ArrayRef< int64_t > strides, int64_t linearIndex)
Given the strides together with a linear index in the dimension space, returns the vector-space offse...
SmallVector< int64_t > getI64SubArray(ArrayAttr arrayAttr, unsigned dropFront=0, unsigned dropBack=0)
Helper that returns a subset of arrayAttr as a vector of int64_t.
mlir::AffineExpr getLinearAffineExpr(ArrayRef< int64_t > basis, mlir::Builder &b)
Computes and returns linearized affine expression w.r.t. basis.
SmallVector< int64_t > invertPermutationVector(ArrayRef< int64_t > permutation)
Helper method to apply to inverse a permutation.
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)
Computes and returns the linearized index of 'offsets' w.r.t. 'basis'.
std::optional< SmallVector< int64_t > > computeShapeRatio(ArrayRef< int64_t > shape, ArrayRef< int64_t > subShape)
Compute and return the multi-dimensional integral ratio of subShape to the trailing dimensions of sha...
bool isPermutationVector(ArrayRef< int64_t > interchange)
Method to check if an interchange vector is a permutation.