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

Quinn Dawkins llvmlistbot at llvm.org
Mon May 20 06:14:33 PDT 2024


================
@@ -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);
----------------
qedawkins wrote:

> Besides this method, I think another possible option is that we can directly RESET corresponding OPERAND of tiled consumer from the result of extractSliceOp(generated by new tiled consumer) to the SOURCE of InsertSliceOp

This solution appears problematic because `getTiledImplementation` is only defined in terms of a tile of the iteration space and the tile of the result that it returns. It makes no requirement that the tiling implementation uses `tensor.extract_slice` to extract the slice of the tensor operand that it wants.

Moreover this IR is UB unless %0 is forwardable to %4 as it is potentially reading uninitialized data.
```
%2, %3 = scf.for iter_args(%dest1, %dest2) {
  %0 = some_producer
  %1 = tensor.insert_slice %0 into %dest1
  %4 = tensor.extract_slice %1
  %5 = tiled_consumer %4
  %6 = tensor.insert %5 into %dest2
  scf.yield %1, %6
}
```
So the only way that I see to guarantee valid tiling is if `getTiledImplementationFromOperandTile` takes the operand tile it is expected to consume as an input.

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


More information about the Mlir-commits mailing list