| //===- TilingInterface.td - Interface for tiling operations *- tablegen -*-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file contains an interface to allow operations to generate a tiled |
| // implementation of themselves. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef MLIR_TILINGINTERFACE |
| #define MLIR_TILINGINTERFACE |
| |
| include "mlir/IR/OpBase.td" |
| |
| def TilingInterface : OpInterface<"TilingInterface"> { |
| let description = [{ |
| Interface for allowing operations to expose information needed to |
| tile them (similar to LinalgOp, but without having access to |
| indexing maps) |
| }]; |
| let cppNamespace = "::mlir"; |
| let methods = [ |
| InterfaceMethod< |
| /*desc=*/[{ |
| Returns a list of iterator types that describe the number of loops. |
| }], |
| /*retType=*/"SmallVector<utils::IteratorType>", |
| /*methodName=*/"getLoopIteratorTypes", |
| /*args=*/(ins), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/"return {};" |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Returns a list of ranges that describe the loop bounds and |
| step for the loops of the operation. |
| }], |
| /*retTy=*/"SmallVector<Range>", |
| /*methodName=*/"getIterationDomain", |
| /*args=*/(ins "OpBuilder &":$b), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/"return {};" |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Method to generate the tiled implementation of an operation. |
| |
| The iteration space of the operation is returned by |
| `getIterationDomain`. The caller provides the information of the |
| tile within this iteration space whose implementation the |
| caller needs. |
| - `offsets` provides the offset of the tile in the coordinate system |
| of the original iteration space, i.e., if an iteration space |
| dimension had non-zero offset, it must be included in the offset |
| provided here (as opposed to zero-based offset "relative" to the |
| iteration space). |
| - `sizes` provides the size of the tile. |
| |
| The method returns the operation that is the tiled |
| implementation. |
| }], |
| /*retType=*/"FailureOr<TilingResult>", |
| /*methodName=*/"getTiledImplementation", |
| /*args=*/(ins |
| "OpBuilder &":$b, |
| "ArrayRef<OpFoldResult> ":$offsets, |
| "ArrayRef<OpFoldResult> ":$sizes), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return {}; |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Method to return the position of the result tile computed by the tiled operation. |
| |
| Specifies what tile of the result of the original tensor is computed |
| by the tiled implementation. Expects the same `offsets` and `sizes` as |
| used to obtain the tiled implementation of the operation. |
| }], |
| /*retType=*/"LogicalResult", |
| /*methodName=*/"getResultTilePosition", |
| /*args=*/(ins |
| "OpBuilder &":$b, |
| "unsigned":$resultNumber, |
| "ArrayRef<OpFoldResult> ":$offsets, |
| "ArrayRef<OpFoldResult> ":$sizes, |
| "SmallVector<OpFoldResult> &":$resultOffsets, |
| "SmallVector<OpFoldResult> &":$resultSizes), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return failure(); |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Method to generate the code that produces a tile of the result. |
| |
| Generates the IR that computes the tile of a result of the |
| operation. The `offsets` and `sizes` describe the tile of |
| the output required. This is different from |
| `getTiledImplementation` which generates the tiled |
| implementation of the operation given a tile of the |
| iteration space. This method generates a tiled |
| implementation of the operation based on the tile of the |
| result required. This method enables fusion by using tile |
| and fuse. The method returns failure if the operation can't be |
| tiled to generate the result tile. In practical terms this |
| implies it cannot be tiled and fused with its consumers. |
| |
| - `offsets` provides the offset of the tile in the coordinate system |
| of the original iteration space, i.e., if an iteration space |
| dimension had non-zero offset, it must be included in the offset |
| provided here (as opposed to zero-based offset "relative" to the |
| iteration space). |
| - `sizes` provides the size of the tile. |
| }], |
| /*retType=*/"FailureOr<TilingResult>", |
| /*methodName=*/"generateResultTileValue", |
| /*args=*/(ins |
| "OpBuilder &":$b, |
| "unsigned":$resultNumber, |
| "ArrayRef<OpFoldResult>":$offsets, |
| "ArrayRef<OpFoldResult>":$sizes), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return failure(); |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Generates the scalar implementation of the operation. |
| |
| Given the list `ivs` that represent points in the iteration space |
| (as specified by `getIterationDomain()`) returns the scalar operations |
| that represent the computation at that point in the iteration space. |
| This method is typically used as the "exit path", i.e. once all |
| transformations are done, this method can be used to lower to scalar |
| code that can then be lowered to LLVM or SPIR-V dialects. |
| }], |
| /*retType=*/"LogicalResult", |
| /*methodName=*/"generateScalarImplementation", |
| /*args=*/(ins |
| "OpBuilder &":$b, |
| "Location ":$loc, |
| "ValueRange ":$ivs), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return failure(); |
| }] |
| > |
| ]; |
| } |
| |
| def PartialReductionOpInterface : OpInterface<"PartialReductionOpInterface"> { |
| let description = [{ |
| Interface for allowing operations to expose information needed to |
| tile reductions using partial reduction followed by merge. This is |
| complementary to TilingInterface to tile reductions. |
| }]; |
| let cppNamespace = "::mlir"; |
| let methods = [ |
| InterfaceMethod< |
| /*desc=*/[{ |
| Method to generate a tensor initalized with the identity value of the |
| operation reduction. The tensor shape is equal to operation result |
| shape with new dimension for each non zero tile size. |
| }], |
| /*retType=*/"FailureOr<Operation*>", |
| /*methodName=*/"generateInitialTensorForPartialReduction", |
| /*args=*/(ins |
| "OpBuilder &":$b, |
| "Location ":$loc, |
| "ArrayRef<OpFoldResult>":$sizes, |
| "ArrayRef<int>":$reductionDim), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return failure(); |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Method to generate a tiled version of the operation where the tiled |
| reduction dimension are converted to parallel dimensions with a size |
| less or equal to the tile size. This is meant to be used with |
| `mergeReductions` method which will combine the partial reductions. |
| }], |
| /*retType=*/"Operation*", |
| /*methodName=*/"tileToPartialReduction", |
| /*args=*/(ins |
| "OpBuilder &":$b, |
| "Location ":$loc, |
| "ValueRange":$init, |
| "ArrayRef<OpFoldResult>":$offsets, |
| "ArrayRef<OpFoldResult>":$sizes, |
| "ArrayRef<int>":$reductionDims), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return nullptr; |
| }] |
| >, |
| InterfaceMethod< |
| /*desc=*/[{ |
| Method to merge partial reductions for an operation that has been |
| tiled along the reduction dimensions. This will only apply the |
| reduction the operation. |
| }], |
| /*retType=*/"Operation*", |
| /*methodName=*/"mergeReductions", |
| /*args=*/(ins |
| "OpBuilder &":$b, |
| "Location ":$loc, |
| "ValueRange":$partialReduce, |
| "ArrayRef<int>":$reductionDim), |
| /*methodBody=*/"", |
| /*defaultImplementation=*/[{ |
| return nullptr; |
| }] |
| > |
| ]; |
| } |
| #endif // MLIR_TILINGINTERFACE |