[mlir][bufferize] Fix missing copy when bufferizing loops

[springerm]

Summary

Using a loop init_arg inside of the loop is not supported. This change adds a pre-processing pass that resolves such IR with copies.

[jreiffers]

Thanks for the quick fix!

mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp

70

You could stop iterating after this. Or consider using llvm::any_of, that way you can get rid of the continue above as well.

You might also get rid of usesInsideRegion completely, if you initialize tensorCopy lazily:

Value tensorCopy;
for (OpOperand &use : ...) {
  if (isProperAncestor && llvm::any_of(...)) {
    if (!tensorCopy) {
      tensorCopy = ...;
    }
    use.set(tensorCopy);
  }
}

Not 100% sure it's nicer though.

[springerm]

address comments

[springerm]

fix typo

mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp

 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
+/// Resolve all operands that are also used inside of repetitive regions of the
+/// same op. Such cases are not fully supported by One-Shot Bufferize.
+///
+/// E.g.:
+/// %r = scf.for ... iter_args(%t = %tensor) -> tensor<?xf32> {
+///   "some_use"(%tensor)
+///   ...
+/// }
+///
+/// Is converted to:
+/// %tensor_copy = bufferization.alloc_tensor copy(%tensor)
+/// %r = scf.for ... iter_args(%t = %tensor) -> tensor<?xf32> {
+///   "some_use"(%tensor_copy)
+///   ...
+/// }
+static void
+resolveUsesInRepetitiveRegions(Operation *op,
+                               const BufferizationOptions &options) {
+  IRRewriter rewriter(op->getContext());
+  AnalysisState state(options);
+
+  // Look for repetitive ops (loops).
+  op->walk([&](RegionBranchOpInterface regionBranchOp) {
+    // Skip non-bufferizable ops.
+    auto bufferizableOp = options.dynCastBufferizableOp(regionBranchOp);
+    if (!bufferizableOp)
+      return WalkResult::advance();
+
+    // Find all operands that are also used inside of a repetitve region of this
+    // op.
+    for (OpOperand &opOperand : regionBranchOp->getOpOperands()) {
+      Value operand = opOperand.get();
+      // Skip non-tensor operands.
+      if (!operand.getType().isa<TensorType>())
+        continue;
+      // Skip operands that do not bufferize to memory writes.
+      if (!bufferizableOp.bufferizesToMemoryWrite(opOperand, state))
+        continue;
+
+      // Gather all uses inside repetitive regions.
+      SmallVector<OpOperand *> usesInsideRegion;
+      for (OpOperand &use : operand.getUses()) {
+        Operation *owner = use.getOwner();
+        if (!regionBranchOp->isProperAncestor(owner))
+          continue;
+        for (Region &r : regionBranchOp->getRegions()) {
+          if (r.findAncestorOpInRegion(*owner) &&
+              regionBranchOp.isRepetitiveRegion(r.getRegionNumber())) {
+            usesInsideRegion.push_back(&use);
+            break;
+          }
+        }
+      }
+      // Nothing to do if the operand is not used inside a repetitive region.
+      if (usesInsideRegion.empty())
+        continue;
+
+      // Insert a tensor copy and replace all uses inside of repetitive regions.
+      rewriter.setInsertionPoint(regionBranchOp);
+      auto tensorCopy = rewriter.create<AllocTensorOp>(
+          regionBranchOp->getLoc(), operand.getType().cast<TensorType>(),
+          /*dynamicSizes=*/ValueRange(),
+          /*copy=*/operand, /*memory_space=*/IntegerAttr());
+      for (OpOperand *use : usesInsideRegion)
+        use->set(tensorCopy);
+    }
+
+    return WalkResult::advance();
+  });
+}
+
 LogicalResult mlir::bufferization::insertTensorCopies(
     Operation *op, const OneShotBufferizationOptions &options) {
+  // Preprocessing: Resolve currently unsupported bufferization cases.
+  resolveUsesInRepetitiveRegions(op, options);
+
   OneShotAnalysisState state(op, options);
   // Run normal One-Shot Bufferize analysis or One-Shot Module Bufferize
   // analysis depending on whether function boundary bufferization is enabled or
   // not.
   if (options.bufferizeFunctionBoundaries) {
     if (failed(analyzeModuleOp(cast<ModuleOp>(op), state)))
       return failure();
   } else {

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

func.func @matmul(
   %c32 = arith.constant 32 : index
   %cst = arith.constant 0.000000e+00 : f32
   %c128 = arith.constant 128 : index
   %c192 = arith.constant 192 : index
   %c8 = arith.constant 8 : index
   %c16 = arith.constant 16 : index
 
   // Hoisted alloc.
-  // CHECK: %[[ALLOC:.*]] = memref.alloc() {alignment = 128 : i64} : memref<8x16xf32>
+  // CHECK: %[[ALLOC:.*]] = memref.alloc() {alignment = 128 : i64} : memref<128x192xf32>
+  // CHECK: memref.copy %[[C]], %[[ALLOC]]
 
   // CHECK: scf.for %[[I:.*]] =
   %0 = scf.for %arg3 = %c0 to %c128 step %c8 iter_args(%arg4 = %C) -> (tensor<128x192xf32>) {
     %1 = tensor.extract_slice %A[%arg3, 0] [8, 256] [1, 1] :
       tensor<128x256xf32> to tensor<8x256xf32>
 
     // CHECK: scf.for %[[J:.*]] =
     %2 = scf.for %arg5 = %c0 to %c192 step %c16 iter_args(%arg6 = %arg4) -> (tensor<128x192xf32>) {
       %3 = tensor.extract_slice %B[0, %arg5] [256, 16] [1, 1] :
         tensor<256x192xf32> to tensor<256x16xf32>
 
-      // %4 does not match an insert_slice, it cannot be bufferized inplace and needs to alloc.
+      // C was already replaced with a copy by preprocessing, so no copy is
+      // needed here.
+      // CHECK: %[[C_SLICE:.*]] = memref.subview %[[ALLOC]]
       %4 = tensor.extract_slice %C[%arg3, %arg5] [8, 16] [1, 1] :
         tensor<128x192xf32> to tensor<8x16xf32>
 
       // linalg.fill is inplace.
-      // CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[ALLOC]] : memref<8x16xf32>)
+      // CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[C_SLICE]]
       %5 = linalg.fill ins(%cst : f32) outs(%4 : tensor<8x16xf32>) -> tensor<8x16xf32>
 
       // CHECK: scf.for %[[K:.*]] =
       %6 = scf.for %arg7 = %c0 to %c256 step %c32 iter_args(%arg8 = %5) -> (tensor<8x16xf32>) {
         %8 = tensor.extract_slice %1[0, %arg7] [8, 32] [1, 1] :
           tensor<8x256xf32> to tensor<8x32xf32>
         %9 = tensor.extract_slice %3[%arg7, 0] [32, 16] [1, 1] :
           tensor<256x16xf32> to tensor<32x16xf32>
 
         // linalg.matmul is inplace as well as the enclosing scf.for.
-        // CHECK: linalg.matmul ins({{.*}} outs(%[[ALLOC]]
+        // CHECK: linalg.matmul ins({{.*}} outs(%[[C_SLICE]]
         %10 = linalg.matmul ins(%8, %9 : tensor<8x32xf32>, tensor<32x16xf32>)
                            outs(%arg8 : tensor<8x16xf32>)
           -> tensor<8x16xf32>
         scf.yield %10 : tensor<8x16xf32>
       }
 
       // insert_slice is inplace but its source comes from an equivalent buffer
       // that is not in place. So we must insert a copy of the small buffer into
       // the bigger buffer.
       // CHECK: %[[T:.*]] = memref.subview %[[C]][%[[I]], %[[J]]] [8, 16] [1, 1]
-      // CHECK: memref.copy %[[ALLOC]], %[[T]]
+      // CHECK: memref.copy %[[C_SLICE]], %[[T]]
       %7 = tensor.insert_slice %6 into %arg6[%arg3, %arg5] [8, 16] [1, 1] :
         tensor<8x16xf32> into tensor<128x192xf32>
 
-      // CHECK: memref.dealloc %[[ALLOC]]
       scf.yield %7 : tensor<128x192xf32>
     }
     scf.yield %2 : tensor<128x192xf32>
   }
+
+  // CHECK: memref.dealloc %[[ALLOC]]
   return %0 : tensor<128x192xf32>
 }
 
 // -----
 
 /// This test just checks the produced IR is valid and does not have dominance
 /// errors in the def-use chains.
 

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

 
 // Note: This bufferizes to inefficient code, but bufferization should not see
 // such IR in the first place. The iter_arg would canonicalize away. This test
 // case is just to ensure that the bufferization generates correct code.
 
 // CHECK-LABEL: func @scf_for_yield_non_equivalent(
 //  CHECK-SAME:     %[[t:.*]]: memref<?xf32
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}})
-//       CHECK:   %[[cloned:.*]] = bufferization.clone %[[alloc]]
-//       CHECK:   memref.dealloc %[[alloc]]
+//       CHECK:   memref.copy %[[t]], %[[alloc]]
+//       CHECK:   %[[cloned:.*]] = bufferization.clone %[[t]]
 //       CHECK:   %[[for:.*]] = scf.for {{.*}} iter_args(%[[iter:.*]] = %[[cloned]])
 //   CHECK-DAG:     memref.dealloc %[[iter]]
 //   CHECK-DAG:     %[[alloc2:.*]] = memref.alloc(%{{.*}})
-//       CHECK:     memref.copy %[[t]], %[[alloc2]]
-//       CHECK:     %[[cloned2:.*]] = bufferization.clone %[[alloc2]]
+//       CHECK:     memref.copy %[[alloc]], %[[alloc2]]
+//       CHECK:     %[[alloc2_casted:.*]] = memref.cast %[[alloc2]]
+//       CHECK:     %[[cloned2:.*]] = bufferization.clone %[[alloc2_casted]]
 //       CHECK:     memref.dealloc %[[alloc2]]
 //       CHECK:     scf.yield %[[cloned2]]
+//       CHECK:   memref.dealloc %[[alloc]]
 //       CHECK:   return %[[for]]
 func.func @scf_for_yield_non_equivalent(
     %t: tensor<?xf32>, %lb : index, %ub : index, %step : index) -> tensor<?xf32> {
   %r = scf.for %i = %lb to %ub step %step iter_args(%a = %t) -> tensor<?xf32> {
     scf.yield %t : tensor<?xf32>
   }
 
   return %r : tensor<?xf32>
%r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
     // Yield tensors in different order.
     scf.yield %ttB, %ttA : tensor<?xf32>, tensor<?xf32>
   }
   // CHECK: }
   %f0 = tensor.extract %r0#0[%step] : tensor<?xf32>
   %f1 = tensor.extract %r0#1[%step] : tensor<?xf32>
   return %f0, %f1: f32, f32
 }
+
+// -----
+
+// CHECK-LABEL: func @scf_for_yield_alias_of_non_equivalent(
+func.func @scf_for_yield_alias_of_non_equivalent(%sz: index) -> tensor<?xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant 5.0 : f32
+
+  // CHECK: %[[generate:.*]] = memref.alloc
+  %0 = tensor.generate %sz {
+  ^bb0(%i: index):
+    tensor.yield %cst : f32
+  } : tensor<?xf32>
+
+  // A copy is inserted because %t is used inside the loop.
+  // CHECK: %[[generate_copy:.*]] = memref.alloc
+  // CHECK: memref.copy %[[generate]], %[[generate_copy]]
+  // CHECK: scf.for
+  %r = scf.for %iv = %c0 to %sz step %c1 iter_args(%t = %0) -> tensor<?xf32> {
+    %iv_sub = arith.subi %iv, %c1 : index
+    // CHECK: memref.subview %[[generate_copy]]
+    %ll = tensor.extract_slice %0[%iv_sub][%sz][1] : tensor<?xf32> to tensor<?xf32>
+    %l = tensor.extract %ll[%c0] : tensor<?xf32>
+    %double = arith.mulf %cst, %l : f32
+    // CHECK: memref.store %{{.*}}, %[[generate]]
+    %s = tensor.insert %double into %t[%iv] : tensor<?xf32>
+    scf.yield %s : tensor<?xf32>
+  }
+  return %r : tensor<?xf32>
+}