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