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 12 #include "mlir/Dialect/Vector/IR/VectorOps.h" 13 #include "mlir/IR/BuiltinAttributes.h" 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_ 166