1//===- TilingInterface.td - Interface for tiling operations *- 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// This file contains an interface to allow operations to generate a tiled 10// implementation of themselves. 11// 12//===----------------------------------------------------------------------===// 13 14#ifndef MLIR_TILINGINTERFACE 15#define MLIR_TILINGINTERFACE 16 17include "mlir/IR/OpBase.td" 18 19def TilingInterface : OpInterface<"TilingInterface"> { 20 let description = [{ 21 This interface allows operations to expose information needed to tile them. 22 23 The intent of this interface is to separate the generation of the loop 24 structure (and constructs used for it) from the information needed from 25 the operation to be able to tile them. As a result an implementation of 26 the tiling algorithm (like `scf::tileUsingSCF`) can generate the inter-tile 27 loop structure, and call into the methods of the interface to be able to 28 tile any operation that implements the interface. 29 30 This interface is also meant to help with "tile and fuse", i.e. the process 31 of fusing a producer with a consumer by 32 a) Tiling the consumer 33 b) Based on the tile of the producer used by the tiled consumer, 34 materialize the tiled implementation of a producer to generate that 35 tile (and use it immediately in the consumer) 36 You could also fuse a consumer with a producer by 37 a) Tiling the producer 38 b) Based on the tile produced, materialize the tiled implementation of 39 a consumer that uses this tile. 40 Note that the tile and fuse does not make any calculation on whether it 41 is "profitable to do this", but simply provides a mechansim to implement 42 the transformation when such a fusion is needed by the caller. 43 44 For any operation to be tilable, an operation has to implement the 45 following two methods (see description below) 46 - `getLoopIteratorTypes` 47 - `getIterationDomain` 48 - `getTiledImplementation` 49 - `getResultTilePosition` 50 51 For an operation to be "tiled and fused" with its (already tiled) consumer, 52 an operation has to implement the following additional method (see 53 description below): 54 - `generateResultTileValue` 55 - `getIterationDomainTileFromResultTile` 56 57 For an operation to be "tiled and fused" with its (already tiled) producer, 58 an operation has to implement the following additional methods (see 59 description below): 60 - `getTiledImplementationFromOperandTile` 61 - `getIterationDomainTileFromOperandTile`. 62 }]; 63 let cppNamespace = "::mlir"; 64 let methods = [ 65 InterfaceMethod< 66 /*desc=*/[{ 67 Returns a list of iterator types that describe the number of loops. 68 }], 69 /*retType=*/"::mlir::SmallVector<::mlir::utils::IteratorType>", 70 /*methodName=*/"getLoopIteratorTypes", 71 /*args=*/(ins), 72 /*methodBody=*/"", 73 /*defaultImplementation=*/"return {};" 74 >, 75 InterfaceMethod< 76 /*desc=*/[{ 77 Returns a list of ranges that describe the loop bounds and 78 step for the loops of the operation. 79 }], 80 /*retTy=*/"::mlir::SmallVector<::mlir::Range>", 81 /*methodName=*/"getIterationDomain", 82 /*args=*/(ins "::mlir::OpBuilder &":$b), 83 /*methodBody=*/"", 84 /*defaultImplementation=*/"return {};" 85 >, 86 InterfaceMethod< 87 /*desc=*/[{ 88 Method to generate the tiled implementation of an operation. 89 90 Given a tile of the iteration space (as returned by 91 `getIterationDomain`), generate in-place the code that represents 92 the computation corresponding to that tile of the iteration space. 93 It is the responsibility of the implementation of this method in 94 the operation to generate the slices of the operands needed for the 95 tiled implementation. 96 - `offsets` provides the offset of the tile in the coordinate system 97 of the original iteration space, i.e., if an iteration space 98 dimension had non-zero offset, it will be included in the offset 99 provided here (as opposed to zero-based offset "relative" to the 100 iteration space). 101 - `sizes` provides the size of the tile. 102 103 The returned `TilingResult` must return for each result of the 104 untiled operation, a `Value` that is the result of the tiled 105 operation. 106 }], 107 /*retType=*/"::mlir::FailureOr<::mlir::TilingResult>", 108 /*methodName=*/"getTiledImplementation", 109 /*args=*/(ins 110 "::mlir::OpBuilder &":$b, 111 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets, 112 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes), 113 /*methodBody=*/"", 114 /*defaultImplementation=*/[{ 115 return {}; 116 }] 117 >, 118 InterfaceMethod< 119 /*desc=*/[{ 120 Method to return the position of the result tile computed by the 121 tiled operation. 122 123 For operations that return a value (typically a value of type 124 `RankedTensorType`), the generated tiled computation has to also 125 recompute a replacement for the results of the original operation. 126 The tiled implementation of the operation returns a tile of the 127 result(s). This methods returns information about what part of the 128 result tensor is computed by the tiled implementation. The manner in 129 which these tiles get put together to get the final result is upto 130 the surrounding loop construct. If an operation has no results, (for 131 example an operation that operates only on memrefs), then this method 132 need not be implemented by the operation. 133 - `resultNumber` is the result number of the original operation 134 being processed. 135 - `offsets` provides the offset of the tile in the coordinate system 136 of the original iteration space, i.e., if an iteration space 137 dimension had non-zero offset, it will be included in the offset 138 provided here (as opposed to zero-based offset "relative" to the 139 iteration space). 140 - `sizes` provides the size of the tile. 141 - `resultOffsets` is the offsets of the tile of the result generated 142 by the tiled implementation (returned by value). 143 - `resultSizes` is the size of the tile of the result generated 144 by the tiled implementation (returned by value). 145 146 Note: It is undefined behaviour if there is overlap between the 147 tiles of the result generated by the tiled implementation. 148 }], 149 /*retType=*/"::llvm::LogicalResult", 150 /*methodName=*/"getResultTilePosition", 151 /*args=*/(ins 152 "::mlir::OpBuilder &":$b, 153 "unsigned":$resultNumber, 154 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets, 155 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes, 156 "::mlir::SmallVector<::mlir::OpFoldResult> &":$resultOffsets, 157 "::mlir::SmallVector<::mlir::OpFoldResult> &":$resultSizes), 158 /*methodBody=*/"", 159 /*defaultImplementation=*/[{ 160 return failure(); 161 }] 162 >, 163 InterfaceMethod< 164 /*desc=*/[{ 165 Method to generate the code that produces a tile of the result. 166 167 This method is required to allow operations to be "tiled and fused" 168 with an (already tiled) consumer. Typically, for two operations with 169 producer -> consumer relation ship, to compute a tile of the 170 consumer a `slice` of the producer is needed. This method allows 171 computing that slice of the producer in-place, thereby "fusing" 172 the operations at tile-granularity. This method is different from 173 `getTiledImplementation`, which produces a tiled implementation 174 for a tile of the iteration space. This method produces a tiled 175 implementation based on the tile of producer required. 176 - `resultNumber` is the result of the producer used by the consumer. 177 - `offsets` is the offset of the slice of the producer result used by 178 the tiled implementation of the consumer. 179 - `sizes` is the size of the slice of the producer result used by the 180 consumer. 181 If fusion of the producer with the consumer is not legal for the 182 operation/result, this method should return failure. 183 184 Note: This method only deals with the mechanism of implementing the 185 fusion. In general the fusion might result in recomputation (based on 186 the way the result is produced by the producer and the access pattern 187 used in the consumer to access). This is upto the caller to handle 188 appropriately. 189 }], 190 /*retType=*/"::mlir::FailureOr<::mlir::TilingResult>", 191 /*methodName=*/"generateResultTileValue", 192 /*args=*/(ins 193 "::mlir::OpBuilder &":$b, 194 "unsigned":$resultNumber, 195 "::mlir::ArrayRef<::mlir::OpFoldResult>":$offsets, 196 "::mlir::ArrayRef<::mlir::OpFoldResult>":$sizes), 197 /*methodBody=*/"", 198 /*defaultImplementation=*/[{ 199 return failure(); 200 }] 201 >, 202 InterfaceMethod< 203 /*desc=*/[{ 204 Method to generate the tiled implementation of an operation that uses 205 exactly a tile of the given operand. 206 207 This method is required to allow operations to be "tiled and fused" 208 with an (already tiled) producer. Given a tile of the producer, this 209 method generates the tile of the consumer that uses exactly this 210 produced tile. In some sense it is the "reverse" of 211 `generateResultTileValue`. 212 - `operandNumber` is the result of the producer used by the consumer. 213 - `offsets` is the offset of the slice of the producer result used by 214 the tiled implementation of the consumer. 215 - `sizes` is the size of the slice of the producer result used by the 216 consumer. 217 If it is illegal to fuse with a producer along the given operand for 218 an operation, the implementation should return a failure. 219 }], 220 /*retType=*/"::mlir::FailureOr<::mlir::TilingResult>", 221 /*methodName=*/"getTiledImplementationFromOperandTile", 222 /*args=*/(ins 223 "::mlir::OpBuilder &":$b, 224 "unsigned":$operandNumber, 225 "::mlir::ArrayRef<::mlir::OpFoldResult>":$offsets, 226 "::mlir::ArrayRef<::mlir::OpFoldResult>":$sizes), 227 /*methodBody=*/"", 228 /*defaultImplementation=*/[{ 229 return failure(); 230 }] 231 >, 232 InterfaceMethod< 233 /*desc=*/[{ 234 Method to return the tile of the iteration domain that uses a given 235 tile of the operand. 236 237 This method is required to allow operations to be "tiled and fused" 238 with an (already tiled) producer. Given a tile of an operand, 239 returns the tile of the iteration space that uses this tile. 240 - `operandNumber` is the result of the producer used by the consumer. 241 - `offsets` is the offset of the slice of the producer result used by 242 the tiled implementation of the consumer. 243 - `sizes` is the size of the slice of the producer result used by the 244 consumer. 245 If it is illegal to fuse with a producer along the given operand for 246 an operation, or if this mapping cannot be computed, the 247 implementation should return a failure. 248 249 Note that unlike the "tile consumer and fuse producer" case, the 250 "tile producer and fuse consumer" requires an additional method to get 251 the iteration tile space that encompasses all uses of the given operand 252 tile. The reason for this is, consider 253 ```mlir 254 %1 = scf.for... { 255 %2 = <tiled_producer_op> 256 %3 = tensor.insert_slice %2 into ... 257 scf.yield %3 258 } 259 %4 = <consumer_op>)(... %1... ) 260 ... <some_op>(... %4 ...) 261 ``` 262 263 when fused this becomes 264 ``` 265 %1 = scf.for... { 266 %2 = <tiled_producer_op> 267 %3 = <tiled_consumer_op>(... %2...) 268 %4 = tensor.insert_slice %3 into ... 269 scf.yield %4 270 } 271 ... <some_op>(... %1 ...) 272 ``` 273 274 i.e, when fusing the consumer, the replacement for the result of the 275 consumer needs to be returned to replace the uses of the consumer. 276 For the tile+fuse algorithm to do this it needs information about 277 which tile of the iteration space encompasses all uses of the tile 278 produced and use that to compute what are the results produced. Note 279 that this iteration space might be the entire iteration space of the 280 operation, or multiple operand tiles might map to intersecting 281 iteration spaces. It is upto the caller to make sure that it is still 282 fusable with producer in this scenario, or it must return a failure. 283 284 Note that this method is only used as a way to implement the 285 transformation. It does not provide guarantees on whether such a 286 transformation is profitable. 287 288 For most cases `getTiledImplementationFromOperandTile` could be a 289 implemented using `getIterationDomainTileFromOperandTile` + 290 `getTiledImplementation` methods. 291 }], 292 /*retType=*/"::llvm::LogicalResult", 293 /*methodName=*/"getIterationDomainTileFromOperandTile", 294 /*args=*/(ins 295 "::mlir::OpBuilder &":$b, 296 "unsigned":$operandNumber, 297 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets, 298 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes, 299 "::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainOffsets, 300 "::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainSizes), 301 /*methodBody=*/"", 302 /*defaultImplementation=*/[{ 303 return failure(); 304 }] 305 >, 306 InterfaceMethod< 307 /*desc=*/[{ 308 Method to return the tile of the iteration domain based 309 on the given tile of the certain result. 310 311 This method is required to allow operations to be "tiled and fused" 312 with an (already tiled) consumer. Given a tile of an result, 313 returns the tile of the iteration space that uses this tile. 314 - `resultNumber` is the result of the producer used by the consumer. 315 - `offsets` is the offset of the slice of the producer result used by 316 the tiled implementation of the consumer. 317 - `sizes` is the size of the slice of the producer result used by the 318 consumer. 319 If fusion of the producer with the consumer is not legal for the 320 result, or if this mapping cannot be computed, the implementation 321 should return a failure. 322 323 For most cases `generateResultTileValue` could be a implemented using 324 `getIterationDomainTileFromResultTile` + `getTiledImplementation` 325 methods. 326 }], 327 /*retType=*/"::llvm::LogicalResult", 328 /*methodName=*/"getIterationDomainTileFromResultTile", 329 /*args=*/(ins 330 "::mlir::OpBuilder &":$b, 331 "unsigned":$resultNumber, 332 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets, 333 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes, 334 "::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainOffsets, 335 "::mlir::SmallVectorImpl<::mlir::OpFoldResult> &":$iterDomainSizes), 336 /*methodBody=*/"", 337 /*defaultImplementation=*/[{ 338 return failure(); 339 }] 340 >, 341 InterfaceMethod< 342 /*desc=*/[{ 343 Generates the scalar implementation of the operation. 344 345 Given the list `ivs` that represent points in the iteration space 346 (as specified by `getIterationDomain()`) returns the scalar operations 347 that represent the computation at that point in the iteration space. 348 This method is typically used as the "exit path", i.e. once all 349 transformations are done, this method can be used to lower to scalar 350 code that can then be lowered to LLVM or SPIR-V dialects. 351 }], 352 /*retType=*/"::llvm::LogicalResult", 353 /*methodName=*/"generateScalarImplementation", 354 /*args=*/(ins 355 "::mlir::OpBuilder &":$b, 356 "::mlir::Location ":$loc, 357 "::mlir::ValueRange ":$ivs), 358 /*methodBody=*/"", 359 /*defaultImplementation=*/[{ 360 return failure(); 361 }] 362 > 363 ]; 364} 365 366def PartialReductionOpInterface : OpInterface<"PartialReductionOpInterface"> { 367 let description = [{ 368 Interface for allowing operations to expose information needed to 369 tile reductions using partial reduction followed by merge. This is 370 complementary to TilingInterface to tile reductions. 371 }]; 372 let cppNamespace = "::mlir"; 373 let methods = [ 374 InterfaceMethod< 375 /*desc=*/[{ 376 Method to generate a tensor initalized with the identity value of the 377 operation reduction. The tensor shape is equal to operation result 378 shape with new dimension for each non zero tile size. 379 }], 380 /*retType=*/"::mlir::FailureOr<SmallVector<Value>>", 381 /*methodName=*/"generateInitialTensorForPartialReduction", 382 /*args=*/(ins 383 "::mlir::OpBuilder &":$b, 384 "Location":$loc, 385 "::mlir::ArrayRef<::mlir::OpFoldResult>":$sizes, 386 "::mlir::ArrayRef<int>":$reductionDim), 387 /*methodBody=*/"", 388 /*defaultImplementation=*/[{ 389 return failure(); 390 }] 391 >, 392 InterfaceMethod< 393 /*desc=*/[{ 394 Method to generate a tiled version of the operation where the tiled 395 reduction dimension are converted to parallel dimensions with a size 396 less or equal to the tile size. This is meant to be used with 397 `mergeReductions` method which will combine the partial reductions. 398 }], 399 /*retType=*/"::mlir::FailureOr<TilingResult>", 400 /*methodName=*/"tileToPartialReduction", 401 /*args=*/(ins 402 "::mlir::OpBuilder &":$b, 403 "Location ":$loc, 404 "ValueRange":$init, 405 "::mlir::ArrayRef<::mlir::OpFoldResult>":$offsets, 406 "::mlir::ArrayRef<::mlir::OpFoldResult>":$sizes, 407 "::mlir::ArrayRef<int>":$reductionDims), 408 /*methodBody=*/"", 409 /*defaultImplementation=*/[{ 410 return failure(); 411 }] 412 >, 413 InterfaceMethod< 414 /*desc=*/[{ 415 Method to merge partial reductions for an operation that has been 416 tiled along the reduction dimensions. This will only apply the 417 reduction the operation. 418 }], 419 /*retType=*/"::mlir::FailureOr<MergeResult>", 420 /*methodName=*/"mergeReductions", 421 /*args=*/(ins 422 "::mlir::OpBuilder &":$b, 423 "Location ":$loc, 424 "ValueRange":$partialReduce, 425 "::mlir::ArrayRef<int>":$reductionDim), 426 /*methodBody=*/"", 427 /*defaultImplementation=*/[{ 428 return failure(); 429 }] 430 >, 431 InterfaceMethod< 432 /*desc=*/[{ 433 Method to return the position of the partial result tile computed by 434 the tiled operation. This is same as 435 TilingInterface:::getResultTilePosition, but determines the result 436 tile position for partial reduction. 437 }], 438 /*retType=*/"::llvm::LogicalResult", 439 /*methodName=*/"getPartialResultTilePosition", 440 /*args=*/(ins 441 "::mlir::OpBuilder &":$b, 442 "unsigned":$resultNumber, 443 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$offsets, 444 "::mlir::ArrayRef<::mlir::OpFoldResult> ":$sizes, 445 "::mlir::SmallVector<::mlir::OpFoldResult> &":$resultOffsets, 446 "::mlir::SmallVector<::mlir::OpFoldResult> &":$resultSizes, 447 "::mlir::ArrayRef<int>":$reductionDims), 448 /*methodBody=*/"", 449 /*defaultImplementation=*/[{ 450 return failure(); 451 }] 452 > 453 ]; 454} 455#endif // MLIR_TILINGINTERFACE 456