1//===- VectorInterfaces.td - Vector interfaces -------------*- tablegen -*-===// 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// Defines the interface for operations on vectors. 10// 11//===----------------------------------------------------------------------===// 12 13#ifndef MLIR_INTERFACES_VECTORINTERFACES 14#define MLIR_INTERFACES_VECTORINTERFACES 15 16include "mlir/IR/OpBase.td" 17 18def VectorUnrollOpInterface : OpInterface<"VectorUnrollOpInterface"> { 19 let description = [{ 20 Encodes properties of an operation on vectors that can be unrolled. 21 }]; 22 let cppNamespace = "::mlir"; 23 24 let methods = [ 25 InterfaceMethod< 26 /*desc=*/[{ 27 Return the shape ratio of unrolling to the target vector shape 28 `targetShape`. Return `std::nullopt` if the op cannot be unrolled to the 29 target vector shape. 30 }], 31 /*retTy=*/"::std::optional<::llvm::SmallVector<int64_t, 4>>", 32 /*methodName=*/"getShapeForUnroll", 33 /*args=*/(ins), 34 /*methodBody=*/"", 35 /*defaultImplementation=*/[{ 36 assert($_op->getNumResults() == 1); 37 auto vt = 38 ::llvm::dyn_cast<::mlir::VectorType>($_op.getResult().getType()); 39 if (!vt) 40 return ::std::nullopt; 41 ::llvm::SmallVector<int64_t, 4> res( 42 vt.getShape().begin(), vt.getShape().end()); 43 return res; 44 }] 45 >, 46 ]; 47} 48 49def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> { 50 let description = [{ 51 Encodes properties of a `vector.transfer_read` or `vector.transfer_write` 52 operation. Vector transfer ops have: 53 54 - A shaped value that the op reads from/writes to: a memref or a tensor. 55 - A vector, either as a result or as an operand. 56 - Indices that describe where the transfer from/to the shaped value starts. 57 - An optional mask. 58 - An optional in_bounds array to indicate transfer dimensions that are 59 guaranteed to be in-bounds. 60 - A permutation map to indicate transposes and broadcasts. 61 62 The "vector rank" is the rank of the vector type. E.g.: 63 ``` 64 // Transfer with shaped value rank 2 and vector (transfer) rank 1. 65 %0 = vector.transfer_read %arg0[%c3, %c3], %f0 66 {permutation_map = affine_map<(d0, d1) -> (d0)>} 67 : memref<?x?xf32>, vector<128xf32> 68 ``` 69 70 The "vector transfer rank" is the number of dimensions that participate in 71 the transfer and broadcasts, and matches the number of results in the 72 permutation map. In most cases, the vector rank matches the vector transfer 73 rank; the only exception is when a vector is flattened as part of the 74 transfer (see `getPermutationMap`). 75 }]; 76 let cppNamespace = "::mlir"; 77 78 let methods = [ 79 InterfaceMethod< 80 /*desc=*/[{ 81 Return the `in_bounds` attribute name. 82 }], 83 /*retTy=*/"::mlir::StringRef", 84 /*methodName=*/"getInBoundsAttrName", 85 /*args=*/(ins) 86 >, 87 InterfaceMethod< 88 /*desc=*/[{ 89 Return the `permutation_map` attribute name. 90 }], 91 /*retTy=*/"::mlir::StringRef", 92 /*methodName=*/"getPermutationMapAttrName", 93 /*args=*/(ins) 94 >, 95 InterfaceMethod< 96 /*desc=*/[{ 97 Return the optional in_bounds attribute that specifies for each vector 98 dimension whether it is in-bounds or not. (Broadcast dimensions are 99 always in-bounds). 100 }], 101 /*retTy=*/"::mlir::ArrayAttr", 102 /*methodName=*/"getInBounds", 103 /*args=*/(ins) 104 >, 105 InterfaceMethod< 106 /*desc=*/[{ 107 Return the memref or ranked tensor operand that this operation operates 108 on. In case of a "read" operation, that's the source from which the 109 operation reads. In case of a "write" operation, that's the destination 110 into which the operation writes. 111 TODO: Change name of operand, which is not accurate for xfer_write. 112 }], 113 /*retTy=*/"::mlir::Value", 114 /*methodName=*/"getSource", 115 /*args=*/(ins) 116 >, 117 InterfaceMethod< 118 /*desc=*/[{ 119 Return the vector that this operation operates on. In case of a "read", 120 that's the vector OpResult. In case of a "write", that's the vector 121 operand value that is written by the op. 122 }], 123 /*retTy=*/"::mlir::Value", 124 /*methodName=*/"getVector", 125 /*args=*/(ins) 126 >, 127 InterfaceMethod< 128 /*desc=*/[{ 129 Return the indices that specify the starting offsets into the source 130 operand. The starting offsets are guaranteed to be in-bounds. 131 }], 132 /*retTy=*/"::mlir::OperandRange", 133 /*methodName=*/"getIndices", 134 /*args=*/(ins) 135 >, 136 InterfaceMethod< 137 /*desc=*/[{ 138 Return the permutation map that describes the mapping of vector 139 dimensions to source dimensions, as well as broadcast dimensions. 140 141 The permutation result has one result per vector transfer dimension. 142 Each result is either a dim expression, indicating the corresponding 143 dimension in the source operand, or a constant "0" expression, 144 indicating a broadcast dimension. 145 146 Note: Nested vector dimensions that are flattened by this op are not 147 accounted for in the permutation map. E.g.: 148 ``` 149 // Vector type has rank 4, but permutation map has only 2 results. That 150 // is because there are only 2 transfer dimensions. 151 %0 = vector.transfer_read %arg1[%c3, %c3], %vf0 152 {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} 153 : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32> 154 ``` 155 }], 156 /*retTy=*/"::mlir::AffineMap", 157 /*methodName=*/"getPermutationMap", 158 /*args=*/(ins) 159 >, 160 InterfaceMethod< 161 /*desc=*/[{ 162 Return the mask operand if the op has a mask. Otherwise, return an 163 empty value. 164 }], 165 /*retTy=*/"Value", 166 /*methodName=*/"getMask", 167 /*args=*/(ins) 168 > 169 ]; 170 171 let extraSharedClassDeclaration = [{ 172 /// Return a vector of all in_bounds values as booleans (one per vector 173 /// transfer dimension). 174 ::llvm::SmallVector<bool> getInBoundsValues() { 175 ::llvm::SmallVector<bool> inBounds; 176 for (int64_t i = 0, e = $_op.getTransferRank(); i < e; ++i) 177 inBounds.push_back($_op.isDimInBounds(i)); 178 return inBounds; 179 } 180 181 /// Return the number of leading shaped dimensions (of the "source" operand) 182 /// that do not participate in the permutation map. 183 unsigned getLeadingShapedRank() { 184 return $_op.getShapedType().getRank() - $_op.getTransferRank(); 185 } 186 187 /// Return the mask type if the op has a mask. Otherwise, return an empty 188 /// VectorType. 189 ::mlir::VectorType getMaskType() { 190 return $_op.getMask() 191 ? ::llvm::cast<::mlir::VectorType>($_op.getMask().getType()) 192 : ::mlir::VectorType(); 193 } 194 195 /// Return the shaped type of the "source" operand value. 196 ::mlir::ShapedType getShapedType() { 197 return ::llvm::cast<::mlir::ShapedType>($_op.getSource().getType()); 198 } 199 200 /// Return the number of dimensions that participate in the permutation map. 201 unsigned getTransferRank() { 202 return $_op.getPermutationMap().getNumResults(); 203 } 204 205 /// Return the type of the vector that this operation operates on. 206 ::mlir::VectorType getVectorType() { 207 return ::llvm::cast<::mlir::VectorType>($_op.getVector().getType()); 208 } 209 210 /// Return "true" if at least one of the vector dimensions is a broadcasted 211 /// dimension. 212 bool hasBroadcastDim() { 213 for (unsigned i = 0, e = $_op.getTransferRank(); i < e; ++i) { 214 if ($_op.isBroadcastDim(i)) 215 return true; 216 } 217 return false; 218 } 219 220 /// Return "true" if at least one of the vector dimensions may be 221 /// out-of-bounds. 222 bool hasOutOfBoundsDim() { 223 for (unsigned idx = 0, e = $_op.getTransferRank(); idx < e; ++idx) 224 if (!$_op.isDimInBounds(idx)) 225 return true; 226 return false; 227 } 228 229 /// Return "true" if the specified vector transfer dimension is a 230 /// broadcasted dimension. 231 bool isBroadcastDim(unsigned dim) { 232 auto expr = $_op.getPermutationMap().getResult(dim); 233 auto constExpr = ::llvm::dyn_cast<::mlir::AffineConstantExpr>(expr); 234 return constExpr && constExpr.getValue() == 0; 235 } 236 237 /// Return "true" if the vector transfer dimension `dim` is in-bounds. 238 /// Return "false" otherwise. 239 bool isDimInBounds(unsigned dim) { 240 auto inBounds = $_op.getInBounds(); 241 return ::llvm::cast<::mlir::BoolAttr>(inBounds[dim]).getValue(); 242 } 243 244 /// Helper function to account for the fact that `permutationMap` results 245 /// and `op.getIndices` sizes may not match and may not be aligned. The 246 /// first `getLeadingShapedRank()` indices may just be indexed and not 247 /// transferred from/into the vector. 248 /// For example: 249 /// ``` 250 /// vector.transfer %0[%i, %j, %k, %c0] 251 /// : memref<?x?x?x?xf32>, vector<2x4xf32> 252 /// ``` 253 /// with `permutation_map = (d0, d1, d2, d3) -> (d2, d3)`. 254 /// Provide a zip function to coiterate on 2 running indices: `resultIdx` 255 /// and `indicesIdx` which accounts for this misalignment. 256 void zipResultAndIndexing( 257 ::llvm::function_ref<void(int64_t, int64_t)> fun) { 258 for (int64_t resultIdx = 0, 259 indicesIdx = $_op.getLeadingShapedRank(), 260 eResult = $_op.getTransferRank(); 261 resultIdx < eResult; 262 ++resultIdx, ++indicesIdx) 263 fun(resultIdx, indicesIdx); 264 } 265 266 /// Return the shape of the hyperrectangular slice within the tensor/memref 267 /// operand that is accessed by the transfer op. 268 /// For example: 269 /// ``` 270 /// vector.transfer %w0[%i, %j, %k] { 271 /// permutation_map = affine_map<(d0, d1, d2) -> (d1, d0, 0)>} : 272 /// : tensor<?x?x?xf32>, vector<4x2x6xf32> 273 /// ``` 274 /// returns a shape [2, 4, 1]. 275 SmallVector<int64_t> getTransferChunkAccessed() { 276 SmallVector<int64_t> dimSizes($_op.getPermutationMap().getNumDims(), 1); 277 for (auto vecDims : llvm::zip($_op.getPermutationMap().getResults(), 278 $_op.getVectorType().getShape())) { 279 AffineExpr dim = std::get<0>(vecDims); 280 int64_t size = std::get<1>(vecDims); 281 // Skip broadcast. 282 if (::llvm::isa<AffineConstantExpr>(dim)) 283 continue; 284 dimSizes[::llvm::cast<AffineDimExpr>(dim).getPosition()] = size; 285 } 286 return dimSizes; 287 } 288 }]; 289} 290 291#endif // MLIR_INTERFACES_VECTORINTERFACES 292