[mlir][linalg][bufferize] Reimplementation of scf.for bufferization

[springerm]

Summary

Instead of modifying the existing scf.for op, create a new op with memref OpOperands/OpResults and delete the old op.

New allocations / other memrefs can now be yielded from the loop. This functionality is deactivated by default and guarded against by AssertDestinationPassingStyle.

This change also introduces replaceOp, which will be utilized by all other bufferize implementations in future commits. Bufferization will then no longer rely on old (pre-bufferize) ops to DCE away. Instead old ops are deleted on the spot. This improves debuggability because there won't be any duplicate ops anymore (bufferized + not-yet-bufferized) when dumping IR during bufferization. It is also less fragile because unbufferized IR can no longer silently "hang around" due to an implementation bug.

[springerm]

update

[nicolasvasilache]

Need to table this for now until we discuss the implications of allowing ToTensor / ToMemref in the parent CL.

[springerm]

rebase

[springerm]

update

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

 /// Note: When reaching an end of the reverse SSA use-def chain, that value
 /// is returned regardless of whether it is a memory write or not.
 Value findLastPrecedingWrite(Value value, const BufferizationOptions &options);
 
 /// Dialect-specific bufferization state. Analysis/bufferization information
 /// that is specific to ops from a certain dialect can be stored in derived
 /// variants of this struct.
 struct DialectBufferizationState {
+  DialectBufferizationState() = default;
+
   virtual ~DialectBufferizationState() = default;
+
+  // Copying state is forbidden. Always pass as reference.
+  DialectBufferizationState(const DialectBufferizationState &) = delete;
 };
 
 /// BufferizationState keeps track of memory buffers and provides a variety of
 /// helper functions for dealing with them. In particular,
 /// `BufferizableOpInterface::bufferize` implementation should utilize the
 /// following helper functions.
 ///
 /// * `createAlloc` / `createDealloc` / `createAllocDeallocPair` creates ops
public:
 
   /// Creates a memref deallocation. The given memref buffer must have been
   /// allocated using `createAlloc`.
   void createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer);
 
   /// Creates a memcpy between two given buffers.
   void createMemCpy(OpBuilder &b, Location loc, Value from, Value to);
 
+  /// Replace an op with replacement values. The op is deleted.
+  void replaceOp(Operation *op, ValueRange values);
+
   /// Map tensor values to memref buffers.
+  // TODO: Deprecated. Remove all uses of this op. Use `replaceOp` instead.
   void mapBuffer(ValueRange tensors, ValueRange buffers);
 
   /// Map a tensor value to a memref buffer.
+  // TODO: Deprecated. Remove all uses of this op. Use `replaceOp` instead.
   void mapBuffer(Value tensor, Value buffer);
 
   /// Lookup the memref buffer that is associated to the given tensor value.
   /// Asserts if no buffer is associated.
   Value lookupBuffer(Value tensor);
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpResult opResult) const;
 
   /// Return `true` if the given value is mapped.
+  // TODO: Deprecated. Remove all uses of this op.
   bool isMapped(Value value) const;
 
   /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
   /// a new buffer and copy over data from the existing buffer if out-of-place
   /// bufferization is necessary.
   Value getResultBuffer(OpResult result);
 
   /// Mark `op` as obsolete, so that it is deleted after bufferization.
+  // TODO: Deprecated. Remove all uses of this op.
   void markOpObsolete(Operation *op);
 
   /// Erase all ops that were marked obsolete.
+  // TODO: Deprecated. Remove all uses of this op.
   void eraseObsoleteOps();
 
   /// Return dialect-specific bufferization state.
   template <typename StateT> StateT &getDialectState(StringRef name) {
     // Create state if it does not exist yet.
     if (!dialectState.count(name))
       dialectState[name] = std::make_unique<StateT>();
     return static_cast<StateT &>(*dialectState[name]);

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

if (!aliasInfo.isInPlace(result)) {
     }
     return resultBuffer;
   }
 
   // Bufferizing in-place. No need to allocate a new buffer.
   return operandBuffer;
 }
 
+void mlir::linalg::comprehensive_bufferize::BufferizationState::replaceOp(
+    Operation *op, ValueRange values) {
+  OpBuilder &b = getBuilder();
+  OpBuilder::InsertionGuard g(b);
+
+  // Replace all OpResults with the given values.
+  for (OpResult opResult : op->getOpResults()) {
+    // Skip OpResult if it has no uses.
+    if (opResult.getUses().empty())
+      continue;
+
+    Value replacement = values[opResult.getResultNumber()];
+    if (opResult.getType().isa<TensorType>()) {
+      // The OpResult is a tensor. Such values are replaced with memrefs during
+      // bufferization.
+      assert((replacement.getType().isa<MemRefType>() ||
+              replacement.getType().isa<UnrankedMemRefType>()) &&
+             "tensor op result should be replaced with a memref value");
+      // The existing uses of the OpResult still expect a tensor. Insert a
+      // ToTensorOp. Throughout bufferization, this ToTensorOp will gradually
+      // loose all of its users and eventually DCE away.
+      setInsertionPointAfter(b, replacement);
+      replacement = b.create<bufferization::ToTensorOp>(replacement.getLoc(),
+                                                        replacement);
+    }
+    opResult.replaceAllUsesWith(replacement);
+  }
+
+  op->erase();
+}
+
 LogicalResult
 mlir::linalg::comprehensive_bufferize::bufferize(Region *region,
                                                  BufferizationState &state) {
   for (Block &block : *region)
     if (failed(bufferize(&block, state)))
       return failure();
   return success();
 }
 
 LogicalResult
 mlir::linalg::comprehensive_bufferize::bufferize(Block *block,
                                                  BufferizationState &state) {
+  // Ops may get deleted during the traversal, so do not iterate over `block`
+  // directly.
+  SmallVector<Operation *> ops;
+  ops.reserve(block->getOperations().size());
   for (Operation &op : *block)
-    if (failed(bufferize(&op, state)))
+    ops.push_back(&op);
+  for (Operation *op : ops)
+    if (failed(bufferize(op, state)))
       return failure();
   return success();
 }
 
 LogicalResult
 mlir::linalg::comprehensive_bufferize::bufferize(Operation *op,
                                                  BufferizationState &state) {
   OpBuilder &b = state.getBuilder();
assert((buffer.getType().isa<MemRefType>() ||
           buffer.getType().isa<UnrankedMemRefType>()) &&
          "expected that tensor is mapped to memref");
   return mapping.map(tensor, buffer);
 }
 
 /// Wrapper for better debugging.
 Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupBuffer(
     Value tensor) {
-  // TODO: if key comes from bbArg, forward.
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
-  Value buffer = mapping.lookupOrNull(tensor);
 
+  // Replace "%t = to_tensor %m" with %m.
+  if (auto toTensorOp = tensor.getDefiningOp<bufferization::ToTensorOp>())
+    return toTensorOp.memref();
+
+  Value buffer = mapping.lookupOrNull(tensor);
   if (!buffer) {
     if (options.allowUnknownOps) {
       // `tensor` was not bufferized yet. This should never happen with
       // bufferizable ops.
       assert(!options.dynCastBufferizableOp(tensor) && "tensor is not mapped");
       // Insert to_memref op.
       OpBuilder b(tensor.getContext());
       setInsertionPointAfter(b, tensor);

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.cpp

 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace bufferization_ext {
 
 // TODO: These ops should implement BufferizableOpInterface directly when moved
 // to the Bufferization dialect.
 
-// TODO: These implementations are conservative and will likely have to be
-// loosened for partial bufferization.
-
 /// ToMemrefOp casts a tensor into a memref. The resulting memref is the memory
 /// location of the incoming tensor once it will be bufferized. In the anlysis,
 /// the incoming tensor is assumed to bufferize to a memory read and to an
 /// inplace memory write, since it is unknown what will happen to the resulting
 /// memref.
+///
+/// Note: ToMemrefOp / ToTensorOp are temporary ops that are inserted at the
+/// bufferization boundary. When bufferization is complete, there should be no
+/// such ops left over. If `allowUnknownOps`, such ops may be part of the
+/// resulting IR, but such IR may no longer be bufferizable by Comprehensive
+/// Bufferize.
 struct ToMemrefOpInterface
     : public BufferizableOpInterface::ExternalModel<ToMemrefOpInterface,
                                                     bufferization::ToMemrefOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     // It is unknown whether the resulting MemRef will be read or not.
     return true;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
+    auto toMemrefOp = cast<bufferization::ToMemrefOp>(op);
+
+    // Fold to_memref(to_tensor(x)) to x.
+    if (auto toTensorOp =
+            toMemrefOp.tensor().getDefiningOp<bufferization::ToTensorOp>()) {
+      toMemrefOp.replaceAllUsesWith(toTensorOp.memref());
+      toMemrefOp.erase();
+      return success();
+    }
+
+    // If a ToMemrefOp's tensor operand has not been bufferized yet, the op
+    // remains unchanged. All IR up to this ToMemrefOp has already been
+    // bufferized, unless there were unknown ops that could be bufferized.
+    if (!state.isMapped(toMemrefOp.tensor())) {
+      assert(state.getOptions().allowUnknownOps &&
+             "expected that tensor is mapped");
+      return success();
+    }
+
+    // If a ToMemrefOp's tensor operand has been bufferized, the op can be
+    // removed.
+    Value memref = state.lookupBuffer(toMemrefOp.tensor());
+    // Do not replace a ToMemrefOp with itself. E.g., when bufferizing a
+    // function body, ToMemrefOps were inserted before starting bufferization of
+    // the function body. Such ToMemrefOps are replaced in a separate step after
+    // the function body has been bufferized.
+    if (toMemrefOp.getResult() != memref)
+      toMemrefOp.replaceAllUsesWith(memref);
+
     return success();
   }
 };
 
 /// ToTensorOp conceptually loads a tensor from a memory location. Such ops do
 /// not lower any further, and they should have disappeared by the time the
 /// input is fully bufferized.
 ///

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp

 //===- SCFInterfaceImpl.cpp - SCF Impl. of BufferizableOpInterface --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace scf_ext {
 
 struct ExecuteRegionOpInterface
     : public BufferizableOpInterface::ExternalModel<ExecuteRegionOpInterface,
bool isWritable(Operation *op, Value value, BufferizationState &state) const {
     // inplace from the perspective of ops nested under:
     //   1. Either the matching iter operand is not bufferized inplace and an
     //      alloc + optional copy makes the bbArg itself inplaceable.
     //   2. Or the matching iter operand is bufferized inplace and bbArg just
     //      bufferizes to that too.
     return true;
   }
 
-  LogicalResult bufferize(Operation *op, OpBuilder &b,
+  LogicalResult bufferize(Operation *op, OpBuilder & /*b*/,
                           BufferizationState &state) const {
     auto forOp = cast<scf::ForOp>(op);
+    Block *oldLoopBody = &forOp.getLoopBody().front();
 
-    // Take a guard before anything else.
-    OpBuilder::InsertionGuard g(b);
+    // Use IRRewriter instead of OpBuilder because it has additional helper
+    // functions.
+    IRRewriter rewriter(op->getContext());
+    rewriter.setInsertionPoint(forOp);
+
+    // Indices of all iter_args that have tensor type. These are the ones that
+    // are bufferized.
+    DenseSet<int64_t> indices;
+    for (const auto &it : llvm::enumerate(forOp.initArgs()))
+      if (it.value().getType().isa<TensorType>())
+        indices.insert(it.index());
+
+    // Given a range of values, apply `func` to those marked in `indices`.
+    // Otherwise, store the unmodified value in the result vector.
+    auto convert = [&](ValueRange values,
+                       std::function<Value(Value, int64_t)> func) {
+      SmallVector<Value> result;
+      for (const auto &it : llvm::enumerate(values)) {
+        size_t idx = it.index();
+        Value val = it.value();
+        result.push_back(indices.contains(idx) ? func(val, idx) : val);
+      }
+      return result;
+    };
 
-    for (OpResult opResult : forOp->getResults()) {
-      if (!opResult.getType().isa<TensorType>())
-        continue;
-      // TODO: Atm we bail on unranked TensorType because we don't know how to
-      // alloc an UnrankedMemRefType + its underlying ranked MemRefType.
-      assert(opResult.getType().isa<RankedTensorType>() &&
-             "unsupported unranked tensor");
+    // Construct a new scf.for op with memref instead of tensor values.
+    SmallVector<Value> initArgs =
+        convert(forOp.initArgs(), [&](Value val, int64_t index) {
+          return state.getResultBuffer(forOp->getOpResult(index));
+        });
+    auto newForOp =
+        rewriter.create<scf::ForOp>(forOp.getLoc(), forOp.lowerBound(),
+                                    forOp.upperBound(), forOp.step(), initArgs);
+    Block *loopBody = &newForOp.getLoopBody().front();
+
+    // Set up new iter_args. The loop body uses tensors, so wrap the (memref)
+    // iter_args of the new loop in ToTensorOps.
+    rewriter.setInsertionPointToStart(loopBody);
+    SmallVector<Value> iterArgs =
+        convert(newForOp.getRegionIterArgs(), [&](Value val, int64_t index) {
+          return rewriter.create<bufferization::ToTensorOp>(val.getLoc(), val);
+        });
+    iterArgs.insert(iterArgs.begin(), newForOp.getInductionVar());
 
-      // TODO: More general: Matching bbArg does not bufferize to a read.
-      Value resultBuffer = state.getResultBuffer(opResult);
-      if (!resultBuffer)
-        return failure();
+    // Erase terminator if present.
+    if (iterArgs.size() == 1)
+      rewriter.eraseOp(loopBody->getTerminator());
+
+    // Move loop body to new loop.
+    rewriter.mergeBlocks(oldLoopBody, loopBody, iterArgs);
+
+    // Update scf.yield of new loop.
+    auto yieldOp = cast<scf::YieldOp>(loopBody->getTerminator());
+    rewriter.setInsertionPoint(yieldOp);
+    SmallVector<Value> yieldValues =
+        convert(yieldOp.results(), [&](Value val, int64_t index) {
+          return rewriter.create<bufferization::ToMemrefOp>(
+              val.getLoc(), initArgs[index].getType(), val);
+        });
+    yieldOp.resultsMutable().assign(yieldValues);
 
-      OpOperand &opOperand = forOp.getOpOperandForResult(opResult);
-      BlockArgument bbArg = forOp.getRegionIterArgForOpOperand(opOperand);
-      state.mapBuffer(bbArg, resultBuffer);
-      state.mapBuffer(opResult, resultBuffer);
-    }
+    // Replace loop results.
+    state.replaceOp(op, newForOp->getResults());
 
     // Bufferize loop body.
-    if (failed(comprehensive_bufferize::bufferize(&forOp.region(), state)))
+    if (failed(comprehensive_bufferize::bufferize(loopBody, state)))
       return failure();
 
-    // Finish bufferizing scf::ForOp.
-    auto yieldOp = cast<scf::YieldOp>(&forOp.region().front().back());
-    for (OpOperand &operand : yieldOp->getOpOperands()) {
-      auto tensorType = operand.get().getType().dyn_cast<TensorType>();
-      if (!tensorType)
-        continue;
-
-      OpOperand &forOperand = forOp.getOpOperandForResult(
-          forOp->getResult(operand.getOperandNumber()));
-      auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
-
-      // Buffers are equivalent so the work is already done and we just yield
-      // the bbArg so that it later canonicalizes away.
-      operand.set(bbArg);
-    }
     return success();
   }
 };
 
 LogicalResult mlir::linalg::comprehensive_bufferize::scf_ext::
     AssertDestinationPassingStyle::run(Operation *op, BufferizationState &state,
                                        BufferizationAliasInfo &aliasInfo,
                                        SmallVector<Operation *> &newOps) {

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-partial.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
 
-// TODO: Bufferize result IR of bufferization.
-// TODO: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
-
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 // RUN: mlir-opt %s -allow-unregistered-dialect -test-comprehensive-function-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
+// RUN: mlir-opt %s -allow-unregistered-dialect -test-comprehensive-function-bufferize="dialect-filter=scf allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-SCF
 
 // CHECK-LABEL: func @use_of_unknown_op_1(
 //  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
 func @use_of_unknown_op_1(%t1: tensor<?xf32> {linalg.inplaceable = true})
     -> vector<5xf32> {
   // ToTensorOp is generated because the function is bufferized and has a
   // memref block argument.
   // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
func @simple_tensor_test(%t1 : tensor<?xf32>, %f : f32) -> tensor<?xf32> {
   // CHECK-TENSOR: %[[casted:.*]] = memref.cast %[[alloc]]
   // CHECK-TENSOR: memref.copy %[[t1_memref]], %[[casted]]
   // CHECK-TENSOR: memref.store %{{.*}}, %[[alloc]]
   %0 = tensor.insert %f into %t1[%c0] : tensor<?xf32>
   // CHECK-TENSOR: %[[casted_tensor:.*]] = bufferization.to_tensor %[[casted]]
   // CHECK-TENSOR: return %[[casted_tensor]]
   return %0 : tensor<?xf32>
 }
+
+// -----
+
+// CHECK-SCF-LABEL: func @simple_scf_for(
+//  CHECK-SCF-SAME:     %[[t1:.*]]: tensor<?xf32>
+func @simple_scf_for(
+    %t1: tensor<?xf32>, %sz: index, %step: index, %f: f32) -> tensor<?xf32> {
+  %c0 = arith.constant 0 : index
+
+  // CHECK-SCF: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
+  // CHECK-SCF: %[[alloc:.*]] = memref.alloc
+  // CHECK-SCF: %[[casted:.*]] = memref.cast %[[alloc]]
+  // CHECK-SCF: memref.copy %[[t1_memref]], %[[casted]]
+  // CHECK-SCF: %[[scf_for:.*]] = scf.for %[[iv:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[arg0:.*]] = %[[casted]]) -> ({{.*}}) {
+  %0 = scf.for %iv = %c0 to %sz step %step iter_args(%arg0 = %t1) -> tensor<?xf32> {
+    // CHECK-SCF: %[[arg0_tensor:.*]] = bufferization.to_tensor %[[arg0]]
+    // CHECK-SCF: %[[insert:.*]] = tensor.insert %{{.*}} into %[[arg0_tensor]]
+    %1 = tensor.insert %f into %arg0[%iv] : tensor<?xf32>
+
+    // CHECK-SCF: %[[insert_memref:.*]] = bufferization.to_memref %[[insert]]
+    // CHECK-SCF: scf.yield %[[insert_memref]]
+    scf.yield %1 : tensor<?xf32>
+  }
+  // CHECK-SCF: }
+
+  // CHECK-SCF: %[[scf_for_tensor:.*]] = bufferization.to_tensor %[[scf_for]]
+  // CHECK-SCF: return %[[scf_for_tensor]]
+  return %0 : tensor<?xf32>
+}

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
 //  CHECK-SAME:   %[[t:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
 func @scf_for_yield_only(%A : tensor<?xf32>,
                          %B : tensor<?xf32> {linalg.inplaceable = true},
                          %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
   //     CHECK:   %[[ALLOC_FOR_A:.*]] = memref.alloc
+  //     CHECK:   %[[CASTED:.*]] = memref.cast %[[ALLOC_FOR_A]]
   //     CHECK:   linalg.copy(%[[A]], %[[ALLOC_FOR_A]])
 
   // The first scf.for remains but just turns into dead code.
   %r0 = scf.for %i = %lb to %ub step %step iter_args(%t = %A) -> (tensor<?xf32>) {
     scf.yield %t : tensor<?xf32>
   }
 
   // The second scf.for remains but just turns into dead code.
   %r1 = scf.for %i = %lb to %ub step %step iter_args(%t = %B) -> (tensor<?xf32>) {
     scf.yield %t : tensor<?xf32>
   }
 
   //     CHECK:   memref.dealloc %[[ALLOC_FOR_A]] : memref<?xf32>
-  //     CHECK:   return %[[ALLOC_FOR_A]] : memref<?xf32>
+  //     CHECK:   return %[[CASTED]] : memref<?xf32, #[[$map_1d_dyn]]>
   return %r0, %r1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // Ensure that the function bufferizes without error. This tests pre-order
 // traversal of scf.for loops during bufferization. No need to check the IR,
 // just want to make sure that it does not crash.

 func @scf_for_with_tensor.insert_slice(
    %A : tensor<?xf32>,
    %B : tensor<?xf32> {linalg.inplaceable = true},
    %C : tensor<4xf32>,
    %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
   //     CHECK:   %[[ALLOC_FOR_A:.*]] = memref.alloc
+  //     CHECK:   %[[CASTED:.*]] = memref.cast %[[ALLOC_FOR_A]]
   //     CHECK:   linalg.copy(%[[A]], %[[ALLOC_FOR_A]])
 
   //     CHECK: %[[svA:.*]] = memref.subview %[[ALLOC_FOR_A]][0] [4] [1]
   //     CHECK: %[[svB:.*]] = memref.subview %[[B]][0] [4] [1]
 
   //     CHECK:   scf.for {{.*}}
   // CHECK-NOT: iter_args
   %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
       -> (tensor<?xf32>, tensor<?xf32>)
   {
     // %ttA bufferizes to direct copy of %BUFFER_CAST_C into %svA
     //     CHECK: linalg.copy(%[[C]], %[[svA]])
     %ttA = tensor.insert_slice %C into %tA[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
     // %ttB bufferizes to direct copy of %BUFFER_CAST_C into %BUFFER_CAST_B
     //     CHECK:   linalg.copy(%[[C]], %[[svB]])
     %ttB = tensor.insert_slice %C into %tB[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
     // CHECK-NOT:   scf.yield
     scf.yield %ttA, %ttB : tensor<?xf32>, tensor<?xf32>
   }
 
   //     CHECK:  memref.dealloc %[[ALLOC_FOR_A]] : memref<?xf32>
-  //     CHECK:  return %[[ALLOC_FOR_A]] : memref<?xf32>
+  //     CHECK:  return %[[CASTED]] : memref<?xf32, #[[$map_1d_dyn]]>
   return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Cross function boundary cases.
 //===----------------------------------------------------------------------===//

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

void TestComprehensiveFunctionBufferize::runOnFunction() {
   BufferizationOptions options;
 
   // Enable InitTensorOp elimination.
   options.addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   // TODO: Find a way to enable this step automatically when bufferizing
   // tensor dialect ops.
   options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
-  options.addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   options.allowReturnMemref = allowReturnMemref;
   options.allowUnknownOps = allowUnknownOps;
   options.testAnalysisOnly = testAnalysisOnly;
   options.analysisFuzzerSeed = analysisFuzzerSeed;
 
   if (dialectFilter.hasValue()) {
     options.dialectFilter.emplace();

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

srcs = [
         "lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp",
     ],
     hdrs = [
         "include/mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h",
     ],
     includes = ["include"],
     deps = [
         ":BufferizableOpInterface",
+        ":BufferizationDialect",
         ":IR",
         ":SCFDialect",
         ":Support",
         "//llvm:Support",
     ],
 )
 
 cc_library(