MLIR  15.0.0git
VectorUtils.h
Go to the documentation of this file.
1 //===- VectorUtils.h - Vector Utilities -------------------------*- 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 #ifndef MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
10 #define MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
11 
14 #include "mlir/Support/LLVM.h"
15 
16 #include "llvm/ADT/DenseMap.h"
17 
18 namespace mlir {
19 
20 // Forward declarations.
21 class AffineApplyOp;
22 class AffineForOp;
23 class AffineMap;
24 class Block;
25 class Location;
26 class OpBuilder;
27 class Operation;
28 class ShapedType;
29 class Value;
30 class VectorType;
31 class VectorTransferOpInterface;
32 
33 namespace vector {
34 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
35 /// the type of `source`.
36 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
37 
38 /// Return the result value of reducing two scalar/vector values with the
39 /// corresponding arith operation.
40 Value makeArithReduction(OpBuilder &b, Location loc, CombiningKind kind,
41  Value v1, Value v2);
42 } // namespace vector
43 
44 /// Return the number of elements of basis, `0` if empty.
45 int64_t computeMaxLinearIndex(ArrayRef<int64_t> basis);
46 
47 /// Given the shape and sizes of a vector, returns the corresponding
48 /// strides for each dimension.
49 /// TODO: needs better doc of how it is used.
50 SmallVector<int64_t, 4> computeStrides(ArrayRef<int64_t> shape,
51  ArrayRef<int64_t> sizes);
52 
53 /// Given the target sizes of a vector, together with vector-space offsets,
54 /// returns the element-space offsets for each dimension.
55 SmallVector<int64_t, 4>
56 computeElementOffsetsFromVectorSliceOffsets(ArrayRef<int64_t> sizes,
57  ArrayRef<int64_t> vectorOffsets);
58 
59 /// Computes and returns the multi-dimensional ratio of `superShape` to
60 /// `subShape`. This is calculated by performing a traversal from minor to major
61 /// dimensions (i.e. in reverse shape order). If integral division is not
62 /// possible, returns None.
63 /// The ArrayRefs are assumed (and enforced) to only contain > 1 values.
64 /// This constraint comes from the fact that they are meant to be used with
65 /// VectorTypes, for which the property holds by construction.
66 ///
67 /// Examples:
68 /// - shapeRatio({3, 4, 5, 8}, {2, 5, 2}) returns {3, 2, 1, 4}
69 /// - shapeRatio({3, 4, 4, 8}, {2, 5, 2}) returns None
70 /// - shapeRatio({1, 2, 10, 32}, {2, 5, 2}) returns {1, 1, 2, 16}
71 Optional<SmallVector<int64_t, 4>> shapeRatio(ArrayRef<int64_t> superShape,
72  ArrayRef<int64_t> subShape);
73 
74 /// Computes and returns the multi-dimensional ratio of the shapes of
75 /// `superVector` to `subVector`. If integral division is not possible, returns
76 /// None.
77 /// Assumes and enforces that the VectorTypes have the same elemental type.
78 Optional<SmallVector<int64_t, 4>> shapeRatio(VectorType superVectorType,
79  VectorType subVectorType);
80 
81 /// Constructs a permutation map of invariant memref indices to vector
82 /// dimension.
83 ///
84 /// If no index is found to be invariant, 0 is added to the permutation_map and
85 /// corresponds to a vector broadcast along that dimension.
86 ///
87 /// The implementation uses the knowledge of the mapping of loops to
88 /// vector dimension. `loopToVectorDim` carries this information as a map with:
89 /// - keys representing "vectorized enclosing loops";
90 /// - values representing the corresponding vector dimension.
91 /// Note that loopToVectorDim is a whole function map from which only enclosing
92 /// loop information is extracted.
93 ///
94 /// Prerequisites: `indices` belong to a vectorizable load or store operation
95 /// (i.e. at most one invariant index along each AffineForOp of
96 /// `loopToVectorDim`). `insertPoint` is the insertion point for the vectorized
97 /// load or store operation.
98 ///
99 /// Example 1:
100 /// The following MLIR snippet:
101 ///
102 /// ```mlir
103 /// affine.for %i3 = 0 to %0 {
104 /// affine.for %i4 = 0 to %1 {
105 /// affine.for %i5 = 0 to %2 {
106 /// %a5 = load %arg0[%i4, %i5, %i3] : memref<?x?x?xf32>
107 /// }}}
108 /// ```
109 ///
110 /// may vectorize with {permutation_map: (d0, d1, d2) -> (d2, d1)} into:
111 ///
112 /// ```mlir
113 /// affine.for %i3 = 0 to %0 step 32 {
114 /// affine.for %i4 = 0 to %1 {
115 /// affine.for %i5 = 0 to %2 step 256 {
116 /// %4 = vector.transfer_read %arg0, %i4, %i5, %i3
117 /// {permutation_map: (d0, d1, d2) -> (d2, d1)} :
118 /// (memref<?x?x?xf32>, index, index) -> vector<32x256xf32>
119 /// }}}
120 /// ```
121 ///
122 /// Meaning that vector.transfer_read will be responsible for reading the slice:
123 /// `%arg0[%i4, %i5:%15+256, %i3:%i3+32]` into vector<32x256xf32>.
124 ///
125 /// Example 2:
126 /// The following MLIR snippet:
127 ///
128 /// ```mlir
129 /// %cst0 = arith.constant 0 : index
130 /// affine.for %i0 = 0 to %0 {
131 /// %a0 = load %arg0[%cst0, %cst0] : memref<?x?xf32>
132 /// }
133 /// ```
134 ///
135 /// may vectorize with {permutation_map: (d0) -> (0)} into:
136 ///
137 /// ```mlir
138 /// affine.for %i0 = 0 to %0 step 128 {
139 /// %3 = vector.transfer_read %arg0, %c0_0, %c0_0
140 /// {permutation_map: (d0, d1) -> (0)} :
141 /// (memref<?x?xf32>, index, index) -> vector<128xf32>
142 /// }
143 /// ````
144 ///
145 /// Meaning that vector.transfer_read will be responsible of reading the slice
146 /// `%arg0[%c0, %c0]` into vector<128xf32> which needs a 1-D vector broadcast.
147 ///
148 AffineMap
149 makePermutationMap(Block *insertPoint, ArrayRef<Value> indices,
150  const DenseMap<Operation *, unsigned> &loopToVectorDim);
151 AffineMap
152 makePermutationMap(Operation *insertPoint, ArrayRef<Value> indices,
153  const DenseMap<Operation *, unsigned> &loopToVectorDim);
154 
155 namespace matcher {
156 
157 /// Matches vector.transfer_read, vector.transfer_write and ops that return a
158 /// vector type that is a multiple of the sub-vector type. This allows passing
159 /// over other smaller vector types in the function and avoids interfering with
160 /// operations on those.
161 /// This is a first approximation, it can easily be extended in the future.
162 /// TODO: this could all be much simpler if we added a bit that a vector type to
163 /// mark that a vector is a strict super-vector but it still does not warrant
164 /// adding even 1 extra bit in the IR for now.
165 bool operatesOnSuperVectorsOf(Operation &op, VectorType subVectorType);
166 
167 } // namespace matcher
168 } // namespace mlir
169 
170 #endif // MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
Include the generated interface declarations.
SmallVector< int64_t, 4 > computeStrides(ArrayRef< int64_t > shape, ArrayRef< int64_t > sizes)
Given the shape and sizes of a vector, returns the corresponding strides for each dimension...
Optional< SmallVector< int64_t, 4 > > shapeRatio(ArrayRef< int64_t > superShape, ArrayRef< int64_t > subShape)
Computes and returns the multi-dimensional ratio of superShape to subShape.
SmallVector< int64_t, 4 > computeElementOffsetsFromVectorSliceOffsets(ArrayRef< int64_t > sizes, ArrayRef< int64_t > vectorOffsets)
Given the target sizes of a vector, together with vector-space offsets, returns the element-space off...
static AffineMap makePermutationMap(ArrayRef< Value > indices, const DenseMap< Operation *, unsigned > &enclosingLoopToVectorDim)
Constructs a permutation map from memref indices to vector dimension.
int64_t computeMaxLinearIndex(ArrayRef< int64_t > basis)
Return the number of elements of basis, 0 if empty.
Definition: VectorUtils.cpp:97
Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim)
Helper function that creates a memref::DimOp or tensor::DimOp depending on the type of source...
Definition: VectorUtils.cpp:37
Value makeArithReduction(OpBuilder &b, Location loc, CombiningKind kind, Value v1, Value v2)
Return the result value of reducing two scalar/vector values with the corresponding arith operation...
Definition: VectorUtils.cpp:46