One-shot-bufferize: fix for inconsistent while arg types in before/after.

[jreiffers]

Summary

Currently, if the before and after regions of a while op have
tensor args in different indices, this leads to a crash.

This moves the pass-through check for args to the handling of the
condition block, since that is where the results are produced, so
it's also where copies must be made.

[springerm]

Nice cleanup, thanks!

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

LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
     // other buffers. I.e., the i-th OpResult may alias with the i-th init_arg;
     // but not with any other OpOperand. If a corresponding OpResult/init_arg
     // pair bufferizes to equivalent buffers, this aliasing requirement is
     // satisfied. Otherwise, we cannot be sure and must yield a new buffer copy.
     // (New buffer copies do not alias with any buffer.)
     OpBuilder::InsertionGuard g(rewriter);
     auto whileOp = cast<scf::WhileOp>(op);
     auto conditionOp = whileOp.getConditionOp();
-    auto yieldOp = whileOp.getYieldOp();
-
-    // Indices of all bbArgs that have tensor type. These are the ones that
-    // are bufferized. The "before" and "after" regions may have different args.
-    DenseSet<int64_t> indicesBefore = getTensorIndices(whileOp.getInits());
-    DenseSet<int64_t> indicesAfter =
-        getTensorIndices(whileOp.getAfterArguments());
 
     // For every yielded value, is the value equivalent to its corresponding
     // bbArg?
     DenseSet<int64_t> equivalentYieldsBefore = getEquivalentBuffers(
         whileOp.getBeforeArguments(), conditionOp.getArgs(), state);
     DenseSet<int64_t> equivalentYieldsAfter = getEquivalentBuffers(
         whileOp.getAfterArguments(), whileOp.getYieldOp().getResults(), state);
 
     // Update "before" region.
     rewriter.setInsertionPoint(conditionOp);
     SmallVector<Value> beforeYieldValues;
     for (int64_t idx = 0;
          idx < static_cast<int64_t>(conditionOp.getArgs().size()); ++idx) {
       Value value = conditionOp.getArgs()[idx];
-      if (!indicesBefore.contains(idx) ||
-          equivalentYieldsBefore.contains(idx)) {
+      if (!value.getType().isa<TensorType>() ||
+          (equivalentYieldsAfter.contains(idx) &&
+              equivalentYieldsBefore.contains(idx))) {
         beforeYieldValues.push_back(value);
         continue;
       }
       FailureOr<Value> alloc =
           allocateTensorForShapedValue(rewriter, conditionOp.getLoc(), value,
                                        /*escape=*/true, state.getOptions());
       if (failed(alloc))
         return failure();
       beforeYieldValues.push_back(*alloc);
     }
     rewriter.updateRootInPlace(conditionOp, [&]() {
       conditionOp.getArgsMutable().assign(beforeYieldValues);
     });
 
-    // Update "after" region.
-    rewriter.setInsertionPoint(yieldOp);
-    SmallVector<Value> afterYieldValues;
-    for (int64_t idx = 0;
-         idx < static_cast<int64_t>(yieldOp.getResults().size()); ++idx) {
-      Value value = yieldOp.getResults()[idx];
-      if (!indicesAfter.contains(idx) || equivalentYieldsAfter.contains(idx)) {
-        afterYieldValues.push_back(value);
-        continue;
-      }
-      FailureOr<Value> alloc =
-          allocateTensorForShapedValue(rewriter, yieldOp.getLoc(), value,
-                                       /*escape=*/true, state.getOptions());
-      if (failed(alloc))
-        return failure();
-      afterYieldValues.push_back(*alloc);
-    }
-    rewriter.updateRootInPlace(yieldOp, [&]() {
-      yieldOp.getResultsMutable().assign(afterYieldValues);
-    });
-
     return success();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto whileOp = cast<scf::WhileOp>(op);
 
     assert(whileOp.getBefore().getBlocks().size() == 1 &&

mlir/test/Dialect/SCF/one-shot-bufferize-allow-return-allocs-no-deallocs.mlir

+// RUN: mlir-opt %s \
+// RUN:     -one-shot-bufferize="allow-return-allocs create-deallocs=0" \
+// RUN:     -split-input-file | \
+// RUN: FileCheck %s --dump-input=always
+
+// A regression test to check that different before and after argument types are
+// bufferized successfully.
+func.func @different_before_after_args() -> tensor<f32> {
+  %true = arith.constant true
+  %cst = arith.constant dense<0.0> : tensor<f32>
+  %0 = scf.while (%arg4 = %true) : (i1) -> (tensor<f32>) {
+    scf.condition(%true) %cst : tensor<f32>
+  } do {
+  ^bb0(%arg4: tensor<f32>):
+    scf.yield %true : i1
+  }
+  return %0 : tensor<f32>
+}
+
+// CHECK-LABEL: @different_before_after_args
 No newline at end of file

mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir

%r0, %r1 = scf.while (%w0 = %A, %w1 = %B)
     // CHECK-DAG: %[[yield0:.*]] = bufferization.alloc_tensor() copy(%[[w1]]) {bufferization.escape = [true]} : tensor<5xi1>
     // CHECK-DAG: %[[yield1:.*]] = bufferization.alloc_tensor() copy(%[[w0]]) {bufferization.escape = [true]} : tensor<5xi1>
     // CHECK: scf.condition(%[[condition]]) %[[yield0]], %[[yield1]]
     scf.condition(%condition) %w1, %w0 : tensor<5xi1>, tensor<5xi1>
   } do {
   ^bb0(%b0: tensor<5xi1>, %b1: tensor<5xi1>):
     // CHECK: } do {
     // CHECK: ^bb0(%[[b0:.*]]: tensor<5xi1>, %[[b1:.*]]: tensor<5xi1>):
-    // CHECK-DAG: %[[yield2:.*]] = bufferization.alloc_tensor() copy(%[[b1]]) {bufferization.escape = [true]} : tensor<5xi1>
-    // CHECK-DAG: %[[yield3:.*]] = bufferization.alloc_tensor() copy(%[[b0]]) {bufferization.escape = [true]} : tensor<5xi1>
-    // CHECK: scf.yield %[[yield2]], %[[yield3]]
+    // CHECK: scf.yield %[[b1]], %[[b0]]
     // CHECK: }
     scf.yield %b1, %b0 : tensor<5xi1>, tensor<5xi1>
   }
 
   return %r0, %r1 : tensor<5xi1>, tensor<5xi1>
 }
 
 // -----

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

 
 // CHECK-LABEL: func @scf_while_non_equiv_condition_and_body(
 //  CHECK-SAME:     %[[arg0:.*]]: memref<5xi1, #{{.*}}>, %[[arg1:.*]]: memref<5xi1, #{{.*}}>
 func.func @scf_while_non_equiv_condition_and_body(%arg0: tensor<5xi1>,
                                                   %arg1: tensor<5xi1>,
                                                   %idx: index)
   -> (tensor<5xi1>, tensor<5xi1>)
 {
-  // CHECK: %[[clone1:.*]] = bufferization.clone %[[arg1]]
-  // CHECK: %[[clone0:.*]] = bufferization.clone %[[arg0]]
+  // CHECK-DAG: %[[clone1:.*]] = bufferization.clone %[[arg1]]
+  // CHECK-DAG: %[[clone0:.*]] = bufferization.clone %[[arg0]]
   // CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[clone0]], %[[w1:.*]] = %[[clone1]]) {{.*}} {
   %r0, %r1 = scf.while (%w0 = %arg0, %w1 = %arg1)
       : (tensor<5xi1>, tensor<5xi1>) -> (tensor<5xi1>, tensor<5xi1>) {
     // CHECK: %[[condition:.*]] = memref.load %[[w0]]
     // CHECK: %[[a1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
     // CHECK: memref.copy %[[w1]], %[[a1]]
     // CHECK: memref.dealloc %[[w1]]
     // CHECK: %[[a0:.*]] = memref.alloc() {{.*}} : memref<5xi1>
     // CHECK: memref.copy %[[w0]], %[[a0]]
     // CHECK: memref.dealloc %[[w0]]
     // CHECK: %[[cloned1:.*]] = bufferization.clone %[[a1]]
     // CHECK: memref.dealloc %[[a1]]
     // CHECK: %[[cloned0:.*]] = bufferization.clone %[[a0]]
     // CHECK: memref.dealloc %[[a0]]
     // CHECK: scf.condition(%[[condition]]) %[[cloned1]], %[[cloned0]]
     %condition = tensor.extract %w0[%idx] : tensor<5xi1>
     scf.condition(%condition) %w1, %w0 : tensor<5xi1>, tensor<5xi1>
   } do {
   ^bb0(%b0: tensor<5xi1>, %b1: tensor<5xi1>):
     // CHECK: } do {
     // CHECK: ^bb0(%[[b0:.*]]: memref<5xi1>, %[[b1:.*]]: memref<5xi1>):
     // CHECK: memref.store %{{.*}}, %[[b0]]
-    // CHECK: %[[a3:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-    // CHECK: memref.copy %[[b1]], %[[a3]]
+    // CHECK: %[[casted1:.*]] = memref.cast %[[b1]]
+    // CHECK: %[[casted0:.*]] = memref.cast %[[b0]]
+    // CHECK: %[[cloned1:.*]] = bufferization.clone %[[casted1]]
     // CHECK: memref.dealloc %[[b1]]
-    // CHECK: %[[a2:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-    // CHECK: memref.copy %[[b0]], %[[a2]]
+    // CHECK: %[[cloned0:.*]] = bufferization.clone %[[casted0]]
     // CHECK: memref.dealloc %[[b0]]
-    // CHECK: %[[casted3:.*]] = memref.cast %[[a3]]
-    // CHECK: %[[casted2:.*]] = memref.cast %[[a2]]
-    // CHECK: %[[cloned2:.*]] = bufferization.clone %[[casted2]]
-    // CHECK: memref.dealloc %[[a2]]
-    // CHECK: %[[cloned3:.*]] = bufferization.clone %[[casted3]]
-    // CHECK: memref.dealloc %[[a3]]
-    // CHECK: scf.yield %[[cloned3]], %[[cloned2]]
+    // CHECK: scf.yield %[[cloned1]], %[[cloned0]]
     // CHECK: }
     %pos = "dummy.some_op"() : () -> (index)
     %val = "dummy.another_op"() : () -> (i1)
     %1 = tensor.insert %val into %b0[%pos] : tensor<5xi1>
     scf.yield %b1, %1 : tensor<5xi1>, tensor<5xi1>
   }
 
   // CHECK: return %[[loop]]#0, %[[loop]]#1
func.func @non_tensor_for_arg(%A : tensor<?xf32> {bufferization.writable = true})
   %c2 = arith.constant 2.0 : f32
   %c10 = arith.constant 10 : index
   %r1:2 = scf.for %i = %c0 to %c10 step %c1 iter_args(%idx = %c1, %t = %A) -> (index, tensor<?xf32>) {
     %t2 = tensor.insert %c2 into %t[%idx] : tensor<?xf32>
     scf.yield %idx, %t2 : index, tensor<?xf32>
   }
   return %r1#1 : tensor<?xf32>
 }
 No newline at end of file