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

[llvmlistbot at llvm.org]

================
@@ -160,6 +215,21 @@ struct LinalgOpTilingInterface
     return success();
   }
 
+  FailureOr<TilingResult> getTiledImplementationFromOperandTile(
+      Operation *op, OpBuilder &b, unsigned operandNumber,
+      ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes) const {
+    SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
+    auto tilingInterfaceOp = cast<TilingInterface>(op);
+    if (failed(tilingInterfaceOp.getIterationDomainTileFromOperandTile(
+            b, operandNumber, offsets, sizes, mappedOffsets, mappedSizes))) {
+      return emitError(
+          op->getLoc(),
+          "unable to obtain the iter domain position of the operation.");
+    }
+    return tilingInterfaceOp.getTiledImplementation(b, mappedOffsets,
+                                                    mappedSizes);
----------------
MaheshRavishankar wrote:

Ok, I think this PR has been put through the review wringer a bit too much. I dont think we should use PRs as a way to bike shed ideas. Maybe some clarification of what the interface is set up to do might help restrict the scope of expansions/modifications to first land this, and then think of how to change the interface to handle it. In general I have strong concerns with the direction of changes suggested in this comment thread. Ill answer those later, but first let me write down what each method in the interface does.

1) `getTiledImplementation` is designed to be as less opinionated as possible. It deals with just tiling of an operation. The operation publishes an iteration space, and given an tile of that iteration space, this method is meant to generate the tiled implementation of that operation. In process it might take slices of the operation/it might not, that is left to the implementation of the operation. This is the most basic method an operation has to implement for it to be tilable, and is potentially the most general way you can achieve tiling. The tiling algorithm is then very simple 
   a) Use the iteration domain to produce the inter-tile iteration space
   b) Then in the inner most loop body, where you know the tile of the iteration space, pass that tile information to this method to generate the loop body.

2) `getResultTilePosition` is also needed for implementing tiling. For a given tile of the iteration space, return the tile of the result that is produced by the tiled implementation. Most tensor tiling algorithms need to stitch back the original tensor value from the tiled operation. This method gives you back information that you need to stitch back the original tensor from the tiled implentation (by inserting the tiled result in the right place of the result tensor)

This just takes care of Tiling.  The following methods take care of fusion

1) `generateResultTileValue` deals with tile and fuse, specifically fusion of tiled consumer with untiled producer. Again it tries to be as less opinionated as possible about the operation semantics. All  this method expects is, given a tile of the result to be computed, generate a tiled implementation of the producer that computes the tile used by the consumer. If the operation cannot do it, it returns a failure. This method might or might not reuse the implementation in `getTiledImplementation`. That depends on op semantics. That is why there isnt a default implementation of this method in the current interface definition.

This is all you need for tile + fuse of producer. Now when we come to tile and fuse of producer, the flip of this is all we should need. IIRC this is what this patch initially started with but through many many rounds of reviews has deflected from initial intent (which is now lost in all the comment threads). All we should need for fusion of the tiled producer with untiled consumer is the following

1) `generateOperandTileValue`, i.e. given a tile of an operand, generate the implementation of the consumer that uses exactly this tile of the operand (I confess the name here is not great, open to naming suggestions). This method again might or might not reuse the implementation in `getTiledImplementation`. That again depends on op semantics which is a black box. If it cant do so, it returns a failure, and the fusion cannot happen. 

I think things get muddled when we try to have a default implementation of the two fusion methods. Default implementation now need to take an opinion on how the tiling of each operation works. I am fine with having a default implementation for the above methods, but IMO they are not core to the interface. They are just convenience methods that reduce implementation burden on operations that implement the interface. If it doesnt fit the semantics of the operation any operation can override this. Some helper methods that could be added

1) `getIterationSpaceTileFromResultTile` - given information of a tile of the result, compute the tile of the iteration space that when used with `getTiledImplementation` generates the required result tile. This can then be used to have a default implementation of `generateResultTileValue`

2) `getIterationSpaceTileFromOperandTile` - given information of a tile of the operand, compute the tile of the iteration space that when used with `getTiledImplementation` generates the tiled implementation that corresponds to the use of exactly the tile of the operand. This can be used to have a default implementation of `generateOperandTileValue`.

I dont think we need to change the `TilingInterface` fundamentally. All the base methods are meant to be as unopinionated as possible so as to allow more operations to be tiled (and fused based on methods the operation implements).

I suggest we first try to land this PR with something that is equivalent of `generateOperandTileValue` and then come back to maybe having default implementation of the different methods. 

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