[Mlir-commits] [mlir] [MLIR][SCF] Add an API to fuse consumer to a producer within scf loop (PR #88712)

[llvmlistbot at llvm.org]

================
@@ -131,6 +152,73 @@ def TilingInterface : OpInterface<"TilingInterface"> {
           return failure();
         }]
       >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Method to generate slices of the operands.
+
+          - `offsets` provides the offset of the tile in the coordinate system
+            of the original coordinate space, i.e., if a dimension from the 
+            coordinate space of the operand has a non-zero offset, it must be
+            included in the offset provided here (as opposed to zero-based offset
+            "relative" to the coordinate space).
+          - `sizes` provides the size of the tile.
+        }],
+        /*retType=*/"SmallVector<Value>",
+        /*methodName=*/"getOperandTilesForIterationDomainTile",
+        /*args=*/(ins
+          "OpBuilder &":$b,
+          "ArrayRef<OpFoldResult>":$offsets,
+          "ArrayRef<OpFoldResult>":$sizes),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return {};
+        }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Method to generate the tiled implementation of an operation from
+          operand tile position.
+
+          Generates the IR that computes the tiled implementation of an
+          operation from operand tile.  The `offsets` and `sizes`
+          describe the tile of the operand 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
+          operand required. This method enables consumer fusion by using
+          tile and fuse. The method returns failure if the operation
+          can't be tiled to generate the operand tile. In practical terms
+          this implies it cannot be tiled and fused with its producers.
+
+          - `offsets` provides the offset of the tile in the coordinate system
+            of the original coordinate space, i.e., if a dimension from the 
+            coordinate space of the operand has a non-zero offset, it must be
+            included in the offset provided here (as opposed to zero-based offset
+            "relative" to the coordinate space).
+          - `sizes` provides the size of the tile.
+        }],
+        /*retType=*/"FailureOr<::mlir::TilingResult>",
+        /*methodName=*/"getTiledImplementationFromOperandTile",
+        /*args=*/(ins
+          "OpBuilder &":$b,
+          "unsigned":$operandNumber,
+          "ArrayRef<OpFoldResult>":$offsets,
+          "ArrayRef<OpFoldResult>":$sizes),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          ::llvm::SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
+          Operation* op = $_op.getOperation();
+          auto tilingInterfaceOp = cast<::mlir::TilingInterface>(op);
+          if (failed(tilingInterfaceOp.getIterationDomainTileFromOperandTile(
+                  b, operandNumber, offsets, sizes, mappedOffsets, mappedSizes))) {
+            return failure();
+          }
+          SmallVector<Value> tiledOperands = tilingInterfaceOp.getOperandTilesForIterationDomainTile(b, mappedOffsets, mappedSizes);
+          tiledOperands[operandNumber] = op->getOperand(operandNumber);
+          return tilingInterfaceOp.getTiledImplementation(b, tiledOperands, mappedOffsets, mappedSizes);
----------------
Yun-Fly wrote:

It seems that you are using mixture fix solution to address issue @qedawkins raised early.
1. In you latest commit, you have already replaced unnecessary `extract_slice` with tiled operand(source of `insert_slice`) [here](https://github.com/llvm/llvm-project/blob/611cb4e750c1975455f30c7865199140ba687fb0/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp#L1434) as @MaheshRavishankar suggested in this [thread](https://github.com/llvm/llvm-project/pull/88712#discussion_r1609384470). This way is also what I preferred in [above  thread](https://github.com/llvm/llvm-project/pull/88712#discussion_r1606369027), which will introduce no change for current `getTiledImplementation` and no impact on downstreams. Furthermore, @MaheshRavishankar use another new cloned `insert_slice` to solve the invalid `extract_slice` semantic after `getTiledImplementation`.

2. However, you also keep another feasible solution here as @qedawkins suggested: add new argument `tiledOperands` for `getTiledImplementation`, before which, you have set`tiledOperands[operandNumber] = op->getOperand(operandNumber);`. However, notice that, you have already replaced operand of `clonedConsumer` by result of cloned `insert_slice` [here](https://github.com/llvm/llvm-project/blob/611cb4e750c1975455f30c7865199140ba687fb0/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp#L1420). So, you are setting original operand rather than expected tiled one here. 

It is up-to you to select which one as final solution in this PR. But they maybe incompatible, in another word, you could only leave only one method left. Otherwise, it will lead to conflict. If you choose second one, suggest to add document to tell caller that only tiled operand is expected.

https://github.com/llvm/llvm-project/pull/88712