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