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"
27 int64_t
linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis);
32 SmallVector<int64_t>
delinearize(ArrayRef<int64_t> strides,
43 ArrayRef<int64_t> v2);
60 std::optional<SmallVector<int64_t>>
71 template <
typename T,
unsigned N>
75 for (
const auto &en :
enumerate(permutation))
76 auxVec[en.index()] = inVec[en.value()];
88 unsigned dropBack = 0);
95 SmallVector<mlir::AffineExpr>
void dropFront(int64_t arr[N], int64_t *res)
Base type for affine expression.
This class is a general helper class for creating context-global objects like types,...
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...
void applyPermutationToVector(SmallVector< T, N > &inVec, ArrayRef< int64_t > permutation)
Apply the permutation defined by permutation to inVec.
bool isPermutationVector(ArrayRef< int64_t > interchange)
Method to check if an interchange vector is a permutation.