1//===-- SubsetOpInterface.td - Tensor Subsets --------------*- 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#ifndef SUBSET_OP_INTERFACE 10#define SUBSET_OP_INTERFACE 11 12include "mlir/IR/OpBase.td" 13 14def SubsetOpInterface : OpInterface<"SubsetOpInterface"> { 15 let description = [{ 16 This interface can be implemented by ops that operate on tensor subsets. A 17 "subset" is a part of a tensor. This interface describes the subset that 18 an implementing op operates on. Only the specified subset may be accessed by 19 the op. 20 21 Subset ops come in two flavours and ops that implement the 22 `SubsetOpInterface` must also implement one of the respective interfaces. 23 - Insertion flavor: Ops that insert a source value into a destination 24 tensor at the specified subset. Such ops return an updated destination 25 tensor and usually implement the `DestinationStyleOpInterface`. Insertion 26 ops must implement the `SubsetInsertionOpInterface`. 27 - Extraction flavor: Ops that extract at a subset. Extraction ops must 28 implement the `SubsetExtractionOpInterface`. 29 30 How the subset is specified is up to the implementing op. E.g.: 31 - `tensor.extract_slice/insert_slice` describe the subset as a 32 hyperrectangular slice. 33 - `tensor.gather/scatter` describe the subset as list of indices. (Not 34 implemented yet.) 35 }]; 36 37 let cppNamespace = "::mlir"; 38 let methods = [ 39 InterfaceMethod< 40 /*desc=*/[{ 41 Return "true" if this op and the given candidate subset op operate on 42 equivalent subsets. Return "false" if the two subsets are disjoint 43 or cannot be proven to be equivalent. 44 45 This interface method does not have to be implemented if 46 `getAccessedHyperrectangularSlice` is implemented. 47 }], 48 /*retType=*/"bool", 49 /*methodName=*/"operatesOnEquivalentSubset", 50 /*args=*/(ins 51 "::mlir::SubsetOpInterface":$candidate, 52 "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn), 53 /*methodBody=*/"", 54 /*defaultImplementation=*/[{ 55 return ::mlir::detail::defaultOperatesOnEquivalentSubset( 56 $_op, candidate, equivalenceFn); 57 }] 58 >, 59 InterfaceMethod< 60 /*desc=*/[{ 61 Return "true" if this op and the given candidate subset op operate on 62 disjoint subsets. Return "false" if the two subsets are equivalent, 63 overlapping or cannot be proven to be disjoint. 64 65 This interface method does not have to be implemented if 66 `getAccessedHyperrectangularSlice` is implemented. 67 }], 68 /*retType=*/"bool", 69 /*methodName=*/"operatesOnDisjointSubset", 70 /*args=*/(ins 71 "::mlir::SubsetOpInterface":$candidate, 72 "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn), 73 /*methodBody=*/"", 74 /*defaultImplementation=*/[{ 75 return ::mlir::detail::defaultOperatesOnDisjointSubset( 76 $_op, candidate, equivalenceFn); 77 }] 78 >, 79 InterfaceMethod< 80 /*desc=*/[{ 81 If this op operates on a hyperrectangular subset, return a 82 description of the subset in terms of offsets, sizes and strides. 83 Otherwise, return "failure". 84 85 This interface method is a convenience method for the most common case 86 of hyperrectangular subset ops. It is optional. If it is implemented, 87 `operatesOnEquivalentSubset` and `operatesOnDisjointSubset` do not 88 have to be implemented. 89 }], 90 /*retType=*/"::mlir::FailureOr<::mlir::HyperrectangularSlice>", 91 /*methodName=*/"getAccessedHyperrectangularSlice", 92 /*args=*/(ins), 93 /*methodBody=*/"", 94 /*defaultImplementation=*/[{ 95 return ::mlir::failure(); 96 }] 97 >, 98 ]; 99 100 let verify = [{ 101 return ::mlir::detail::verifySubsetOpInterface( 102 ::mlir::cast<::mlir::SubsetOpInterface>($_op)); 103 }]; 104 105 let extraClassDeclaration = [{ 106 /// Return the container that this operation is operating on. In case of an 107 /// extraction op, the container is the source tensor. In case of an 108 /// insertion op, the container is the destination tensor. 109 Value getTensorContainer() { 110 return ::mlir::detail::getTensorContainer(getOperation()); 111 } 112 }]; 113} 114 115def SubsetExtractionOpInterface 116 : OpInterface<"SubsetExtractionOpInterface", [SubsetOpInterface]> { 117 let description = [{ 118 This interface can be implemented by ops that extract a value from 119 a source tensor at a specified subset. The elements in the source tensor 120 that are read by this extraction are called "subset". 121 122 Extraction ops must have a single result value. 123 }]; 124 125 let cppNamespace = "::mlir"; 126 let methods = [ 127 InterfaceMethod< 128 /*desc=*/[{ 129 Return the source tensor operand. 130 }], 131 /*retType=*/"::mlir::OpOperand &", 132 /*methodName=*/"getSourceOperand", 133 /*args=*/(ins) 134 >, 135 ]; 136 137 let verify = [{ 138 return ::mlir::detail::verifySubsetExtractionOpInterface( 139 ::mlir::cast<::mlir::SubsetExtractionOpInterface>($_op)); 140 }]; 141 142 let extraClassDeclaration = [{ 143 /// Return the single result of this op. 144 ::mlir::Value getResult() { 145 return getOperation()->getResult(0); 146 } 147 }]; 148} 149 150def SubsetInsertionOpInterface 151 : OpInterface<"SubsetInsertionOpInterface", [SubsetOpInterface]> { 152 let description = [{ 153 This interface can be implemented by ops that insert a source value into 154 a destination tensor at a specified subset. The elements in the destination 155 tensor that are overwritten by this insertion are called "subset". The 156 updated destination tensor is returned. 157 158 This interface provides helper methods for efficient bufferization of 159 subset-based tensor IR. Tensor subsets can bufferize to buffer "views"/ 160 "aliases" (in contrast to one or multiple less efficient buffer allocation). 161 162 This interface is queried by One-Shot Bufferize to detect cases where a 163 seeming read-after-write is not actually a conflict because the respective 164 ops are operating on equivalent subsets. More details can be found in the 165 documentation of One-Shot Analysis (see `areNonConflictingSubsets`). 166 }]; 167 168 let cppNamespace = "::mlir"; 169 let methods = [ 170 InterfaceMethod< 171 /*desc=*/[{ 172 Return the source operand. The source operand can have any type. 173 }], 174 /*retType=*/"::mlir::OpOperand &", 175 /*methodName=*/"getSourceOperand", 176 /*args=*/(ins) 177 >, 178 InterfaceMethod< 179 /*desc=*/[{ 180 Return the destination operand. The destination operand must be a 181 tensor. 182 183 This function does not have to be implemented for destination style 184 ops that have exactly one "init" operand. 185 }], 186 /*retType=*/"::mlir::OpOperand &", 187 /*methodName=*/"getDestinationOperand", 188 /*args=*/(ins), 189 /*methodBody=*/"", 190 /*defaultImplementation=*/[{ 191 return ::mlir::detail::defaultGetDestinationOperand( 192 $_op.getOperation()); 193 }] 194 >, 195 InterfaceMethod< 196 /*desc=*/[{ 197 Return the updated destination result. 198 199 This function does not have to be implemented for destination style 200 ops. 201 }], 202 /*retType=*/"::mlir::OpResult", 203 /*methodName=*/"getUpdatedDestination", 204 /*args=*/(ins), 205 /*methodBody=*/"", 206 /*defaultImplementation=*/[{ 207 return ::mlir::detail::defaultGetUpdatedDestination( 208 $_op.getOperation()); 209 }] 210 >, 211 InterfaceMethod< 212 /*desc=*/[{ 213 Return "true" if this operation inserts into a subset that is 214 equivalent to the subset defined by `candidate`. 215 216 Two subsets are "equivalent" and "same" if they can bufferize to the 217 same buffer views/aliases. If they are "equivalent", the tensor IR 218 may be expressed in terms of different SSA values (but they could 219 bufferize to MemRef SSA values that can CSE without breaking 220 correctness). `equivalenceFn` should return "true" if the two given 221 values are equivalent. 222 223 Example: 224 ``` 225 // The subset of the SubsetInsertionOpInterface op %1 is equivalent to 226 // the subset defined by %2 (but not "same"): 227 %0 = arith.select %c, %t, %t : tensor<?xf32> 228 %1 = tensor.insert_slice %x into %0[0][5][1] 229 : tensor<5xf32> into tensor<?xf32> 230 %2 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32> 231 232 // The subset of the SubsetInsertionOpInterface op %1 is equivalent to 233 // and "same" as the subset defined by %2. 234 %1 = tensor.insert_slice %x into %t[0][5][1] 235 : tensor<5xf32> into tensor<?xf32> 236 %2 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32> 237 ``` 238 239 Note: By default, this function calls 240 `SubsetOpInterface::operatesOnEquivalentSubset`. 241 }], 242 /*retType=*/"bool", 243 /*methodName=*/"isEquivalentSubset", 244 /*args=*/(ins 245 "::mlir::Value":$candidate, 246 "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn), 247 /*methodBody=*/"", 248 /*defaultImplementation=*/[{ 249 return ::mlir::detail::defaultIsEquivalentSubset( 250 $_op.getOperation(), candidate, equivalenceFn); 251 }] 252 >, 253 InterfaceMethod< 254 /*desc=*/[{ 255 Return the subset of the destination tensor that this operation 256 inserts into. 257 258 Example: 259 ``` 260 // SubsetOpInterface op: 261 %0 = tensor.insert_slice %t0 into %t1[%pos][5][1] 262 : tensor<5xf32> into tensor<?xf32> 263 // Subset (built by this function): 264 %1 = tensor.extract_slice %t1[%pos][5][1] 265 : tensor<?xf32> to tensor<5xf32> 266 ``` 267 268 Note: Implementations do not necessarily have to build new IR. They 269 may return existing SSA values. 270 }], 271 /*retType=*/"::mlir::Value", 272 /*methodName=*/"buildSubsetExtraction", 273 /*args=*/(ins "::mlir::OpBuilder &":$builder, "Location":$loc) 274 >, 275 InterfaceMethod< 276 /*desc=*/[{ 277 Return all SSA values that are needed (i.e., must be in scope) at the 278 insertion of the builder when calling `buildSubsetExtraction`. Users 279 of `buildSubsetExtraction` can use this helper method to find a 280 suitable insertion point. 281 282 Example: The SSA values needed to build the subset in the example of 283 `buildSubsetExtraction` are %t1 and %pos. 284 }], 285 /*retType=*/"::llvm::SmallVector<::mlir::Value>", 286 /*methodName=*/"getValuesNeededToBuildSubsetExtraction", 287 /*args=*/(ins) 288 >, 289 ]; 290 291 let extraClassDeclaration = [{ 292 /// Return "true" if this operation inserts into the same subset as defined 293 /// by `candidate`. 294 /// 295 /// Note: This function is useful outside of bufferization, where no tensor 296 /// equivalence information is available. 297 bool isSameSubset(OpResult candidate) { 298 auto subsetOp = cast<::mlir::SubsetInsertionOpInterface>( 299 getOperation()); 300 return subsetOp.isEquivalentSubset( 301 candidate, [](Value v1, Value v2) { return v1 == v2; }); 302 } 303 }]; 304} 305 306#endif // SUBSET_OP_INTERFACE 307