[MLIR][SCF] Enable better bufferization for `TileConsumerAndFuseProducersUsingSCFForOp`

[chelini]

Summary

Replace iterators of the outermost loop with region arguments of the innermost
one. The changes avoid later bufferization passes to insert allocation within
the body of the innermost loop.

[chelini]

Small change.

[mravishankar]

Thanks! Just a minor suggestion.

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

360

Couple of suggestions here

1. We dont need to walk the body and just replace tensor.extract_slice. It should be OK to just do

Block *block = &(*innerFor.getRegion().begin());
for (auto it : llvm::zip(outerFor.getIterOperands(), innerFor.getRegionIterArgs()) {
  Value source = std::get<0>(it);
  Value target = std::get<1>(it);
  source.replaceAllUsesWithIf(target, [&](OpOperand &use) { return use.getOwner()->getBlock() == block; });
}

2. Not relevant here, but I think updateRootInPlace is really for the root of the pattern.

[chelini]

Update as suggested.

[chelini]

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

360

Updated. Nice solution :)

[mravishankar]

Awesome!

[chelini]

Currently investigating why the test fails. Seems related with: https://github.com/llvm/llvm-project/commit/485190df95f98c51c3f4a4ab4db96127cdc9ce78

[chelini]

Rebase on main

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

if (!loopOp)
       return llvm::None;
     v = loopOp.getOpOperandForRegionIterArg(blockArg).get();
   }
   if (!isa_and_nonnull<TilingInterface>(v.getDefiningOp()))
     return llvm::None;
   return v.cast<OpResult>();
 }
 
+// Replace iter args of the outer most loop with region args of the inner most
+// one.
+static void replaceIterArgs(scf::ForOp outerFor, scf::ForOp innerFor,
+                            PatternRewriter &rewriter) {
+  assert(outerFor.getNumIterOperands() == innerFor.getNumIterOperands() &&
+         "expect same number of iter args");
+  Block *block = &(*innerFor.getRegion().begin());
+  for (auto it :
+       llvm::zip(outerFor.getIterOperands(), innerFor.getRegionIterArgs())) {
+    Value source = std::get<0>(it);
+    Value target = std::get<1>(it);
+    source.replaceUsesWithIf(target, [&](OpOperand &use) {
+      return use.getOwner()->getBlock() == block;
+    });
+  }
+}
+
 FailureOr<scf::SCFTileAndFuseResult>
 scf::TileConsumerAndFuseProducersUsingSCFForOp::returningMatchAndRewrite(
     TilingInterface op, PatternRewriter &rewriter) const {
   // This transformation is only valid for ops that return values (i.e. not
   // valid to use with operations that have memref operands).
   if (!op->getNumResults()) {
     return rewriter.notifyMatchFailure(
         op, "invalid pattern for op with no results");
for (OpOperand &uses : unfusedProducerOp->getUses()) {
       if (uses.getOwner() == outerMostTiledLoop.getOperation()) {
         unsigned resultNumber = uses.get().cast<OpResult>().getResultNumber();
         unsigned operandNumber = uses.getOperandNumber();
         outerMostTiledLoop->setOperand(
             operandNumber, unfusedProducerOpDestValues[resultNumber]);
       }
     }
   }
+  replaceIterArgs(tileAndFuseResult.loops.front(),
+                  tileAndFuseResult.loops.back(), rewriter);
   return tileAndFuseResult;
 }

mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-interface.mlir

 // CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
 //      CHECK:   %[[INIT:.+]] = linalg.init_tensor
 //      CHECK:   scf.for %[[IV0:[a-zA-Z0-9]+]] =
 // CHECK-SAME:       iter_args(%[[ITERARG0:.+]] = %[[INIT]])
 //      CHECK:     scf.for %[[IV1:[a-zA-Z0-9]+]] =
 // CHECK-SAME:         iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]])
 //  CHECK-DAG:       %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
 //  CHECK-DAG:       %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
-//  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT]][%[[IV0]], %[[IV1]]]
+//  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]]
 //      CHECK:       %[[FILL_TILE:.+]] = linalg.fill
 // CHECK-SAME:           outs(%[[INIT_TILE]] :
 //      CHECK:       %[[GEMM_TILE:.+]] = linalg.matmul
 // CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
 // CHECK-SAME:           outs(%[[FILL_TILE]] :
 //      CHECK:       %[[INSERT:.+]] = tensor.insert_slice %[[GEMM_TILE]] into %[[ITERARG1]][%[[IV0]], %[[IV1]]]
 //      CHECK        scf.yield %[[INSERT]]
 

 // CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32>)
 //      CHECK:   %[[INIT:.+]] = linalg.init_tensor
 //      CHECK:   scf.for %[[IV0:[a-zA-Z0-9]+]] =
 // CHECK-SAME:       iter_args(%[[ITERARG0:.+]] = %[[INIT]])
 //      CHECK:     scf.for %[[IV1:[a-zA-Z0-9]+]] =
 // CHECK-SAME:         iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]])
 //  CHECK-DAG:       %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
 //  CHECK-DAG:       %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
-//  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT]][%[[IV0]], %[[IV1]]]
+//  CHECK-DAG:       %[[INIT_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]]
 //      CHECK:       %[[FILL_TILE:.+]] = linalg.fill
 // CHECK-SAME:           outs(%[[INIT_TILE]] :
 //      CHECK:       %[[GEMM_TILE:.+]] = linalg.matmul
 // CHECK-SAME:           ins(%[[LHS_TILE]], %[[RHS_TILE]] :
 // CHECK-SAME:           outs(%[[FILL_TILE]] :
 //  CHECK-DAG:       %[[BIAS_TILE:.+]] = tensor.extract_slice %[[ARG2]][%[[IV1]]]
 //  CHECK-DAG:       %[[OUTS_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]]
 //      CHECK:       %[[GENERIC_TILE:.+]] = linalg.generic

 //  CHECK-DAG:     %[[RHS0_TILE:.+]] = tensor.extract_slice %[[RHS0]][0, 0]
 //  CHECK-DAG:     %[[INIT0_TILE:.+]] = tensor.extract_slice %[[INIT0]][%[[IV]], 0]
 //      CHECK:     %[[FILL0_TILE:.+]] = linalg.fill
 // CHECK-SAME:         outs(%[[INIT0_TILE]] :
 //      CHECK:     %[[GEMM0_TILE:.+]] = linalg.matmul
 // CHECK-SAME:         ins(%[[LHS0_TILE]], %[[RHS0_TILE]] :
 // CHECK-SAME:         outs(%[[FILL0_TILE]] :
 //  CHECK-DAG:     %[[RHS1_TILE:.+]] = tensor.extract_slice %[[RHS1]][0, 0]
-//  CHECK-DAG:     %[[INIT1_TILE:.+]] = tensor.extract_slice %[[INIT1]][%[[IV]], 0]
+//  CHECK-DAG:     %[[INIT1_TILE:.+]] = tensor.extract_slice %[[ITERARG]][%[[IV]], 0]
 //      CHECK:     %[[FILL1_TILE:.+]] = linalg.fill
 // CHECK-SAME:         outs(%[[INIT1_TILE]] :
 //      CHECK:     %[[GEMM1_TILE:.+]] = linalg.matmul
 // CHECK-SAME:         ins(%[[GEMM0_TILE]], %[[RHS1_TILE]] :
 // CHECK-SAME:         outs(%[[FILL1_TILE]] :
 //      CHECK:     %[[INSERT:.+]] = tensor.insert_slice %[[GEMM1_TILE]] into %[[ITERARG]][%[[IV]], 0]
 //      CHECK      scf.yield %[[INSERT]]