MLIR  17.0.0git
IndexingUtils.h
Go to the documentation of this file.
1 //===- IndexingUtils.h - Helpers related to index computations --*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This header file defines utilities and common canonicalization patterns for
10 // reshape operations.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_DIALECT_UTILS_INDEXINGUTILS_H
15 #define MLIR_DIALECT_UTILS_INDEXINGUTILS_H
16 
17 #include "mlir/IR/Builders.h"
18 #include "mlir/Support/LLVM.h"
19 #include "llvm/ADT/ArrayRef.h"
20 #include "llvm/ADT/SmallVector.h"
21 #include <optional>
22 
23 namespace mlir {
24 class ArrayAttr;
25 
26 /// Computes and returns the linearized index of 'offsets' w.r.t. 'basis'.
27 int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis);
28 
29 /// Given the strides together with a linear index in the dimension
30 /// space, returns the vector-space offsets in each dimension for a
31 /// de-linearized index.
32 SmallVector<int64_t> delinearize(ArrayRef<int64_t> strides,
33  int64_t linearIndex);
34 
35 /// Given a set of sizes, compute and return the strides (i.e. the number of
36 /// linear incides to skip along the (k-1) most minor dimensions to get the next
37 /// k-slice). This is also the basis that one can use to linearize an n-D offset
38 /// confined to `[0 .. sizes]`.
39 SmallVector<int64_t> computeStrides(ArrayRef<int64_t> sizes);
40 
41 /// Return a vector containing llvm::zip of v1 and v2 multiplied elementwise.
42 SmallVector<int64_t> computeElementwiseMul(ArrayRef<int64_t> v1,
43  ArrayRef<int64_t> v2);
44 
45 /// Compute and return the multi-dimensional integral ratio of `subShape` to
46 /// the trailing dimensions of `shape`. This represents how many times
47 /// `subShape` fits within `shape`.
48 /// If integral division is not possible, return std::nullopt.
49 /// The trailing `subShape.size()` entries of both shapes are assumed (and
50 /// enforced) to only contain noonnegative values.
51 ///
52 /// Examples:
53 /// - shapeRatio({3, 5, 8}, {2, 5, 2}) returns {3, 2, 1}.
54 /// - shapeRatio({3, 8}, {2, 5, 2}) returns std::nullopt (subshape has higher
55 /// rank).
56 /// - shapeRatio({42, 2, 10, 32}, {2, 5, 2}) returns {42, 1, 2, 16} which is
57 /// derived as {42(leading shape dim), 2/2, 10/5, 32/2}.
58 /// - shapeRatio({42, 2, 11, 32}, {2, 5, 2}) returns std::nullopt which is
59 /// derived as {42(leading shape dim), 2/2, 11/5(not divisible), 32/2}.
60 std::optional<SmallVector<int64_t>>
61 computeShapeRatio(ArrayRef<int64_t> shape, ArrayRef<int64_t> subShape);
62 
63 /// Return the number of elements of basis (i.e. the max linear index).
64 /// Return `0` if `basis` is empty.
65 int64_t computeMaxLinearIndex(ArrayRef<int64_t> basis);
66 
67 /// Apply the permutation defined by `permutation` to `inVec`.
68 /// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
69 /// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation vector
70 /// `permutation = [2, 0, 1]`, this function leaves `inVec = ['c', 'a', 'b']`.
71 template <typename T, unsigned N>
73  ArrayRef<int64_t> permutation) {
74  SmallVector<T, N> auxVec(inVec.size());
75  for (const auto &en : enumerate(permutation))
76  auxVec[en.index()] = inVec[en.value()];
77  inVec = auxVec;
78 }
79 
80 /// Helper method to apply to inverse a permutation.
81 SmallVector<int64_t> invertPermutationVector(ArrayRef<int64_t> permutation);
82 
83 /// Method to check if an interchange vector is a permutation.
84 bool isPermutationVector(ArrayRef<int64_t> interchange);
85 
86 /// Helper that returns a subset of `arrayAttr` as a vector of int64_t.
87 SmallVector<int64_t> getI64SubArray(ArrayAttr arrayAttr, unsigned dropFront = 0,
88  unsigned dropBack = 0);
89 
90 /// Computes and returns linearized affine expression w.r.t. `basis`.
91 mlir::AffineExpr getLinearAffineExpr(ArrayRef<int64_t> basis, mlir::Builder &b);
92 
93 /// Given the strides in the dimension space, returns the affine expressions for
94 /// vector-space offsets in each dimension for a de-linearized index.
95 SmallVector<mlir::AffineExpr>
96 getDelinearizedAffineExpr(ArrayRef<int64_t> strides, mlir::Builder &b);
97 
98 } // namespace mlir
99 
100 #endif // MLIR_DIALECT_UTILS_INDEXINGUTILS_H
void dropFront(int64_t arr[N], int64_t *res)
Definition: CRunnerUtils.h:118
Base type for affine expression.
Definition: AffineExpr.h:68
This class is a general helper class for creating context-global objects like types,...
Definition: Builders.h:50
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:223
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.
Definition: IndexingUtils.h:72
bool isPermutationVector(ArrayRef< int64_t > interchange)
Method to check if an interchange vector is a permutation.