[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/SCFInterfaceImpl.h

 namespace mlir {
 
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace scf_ext {
 
+/// Equivalence analysis for scf.for. Raise an error if iter_args are not
+/// equivalent to their corresponding loop yield values.
+struct AssertDestinationPassingStyle : public PostAnalysisStep {
+  LogicalResult run(Operation *op, BufferizationState &state,
+                    SmallVector<Operation *> &newOps) override;
+};
+
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace scf_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_SCF_INTERFACE_IMPL_H

mlir/include/mlir/Dialect/Linalg/Passes.h

 /// analysis to determine which op operands and results may be bufferized in the
 /// same buffers.
 std::unique_ptr<Pass> createLinalgComprehensiveFunctionBufferizePass();
 
 /// This pass is an extension of the comprehensive function bufferization. It
 /// performs a function call analysis and bufferizes an entire module.
 std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass();
 
+/// This is a pass that bufferizes only SCF dialect ops.
+std::unique_ptr<Pass> createLinalgComprehensiveSCFBufferizePass();
+
 /// This is a pass that bufferizes only tensor dialect ops.
 std::unique_ptr<Pass> createLinalgComprehensiveTensorBufferizePass();
 
 /// This is a pass that bufferizes only vector dialect ops.
 std::unique_ptr<Pass> createLinalgComprehensiveVectorBufferizePass();
 
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.

mlir/include/mlir/Dialect/Linalg/Passes.td

Option<"useAlloca", "use-alloca", "bool",
            "Use stack allocations for memrefs (for testing purposes only)">,
     Option<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
            /*default=*/"0",
            "Analyze ops in random order with a given seed (fuzzer)">
   ];
   let constructor = "mlir::createLinalgComprehensiveModuleBufferizePass()";
 }
 
+def LinalgComprehensiveSCFBufferize :
+    FunctionPass<"linalg-comprehensive-scf-bufferize"> {
+  let summary = "Bufferize (tensor into memref) SCF dialect ops.";
+  let description = [{
+    This pass bufferizes only SCF dialect ops.
+    TODO: This will be moved to the bufferization dialect in the future.
+  }];
+  let constructor = "mlir::createLinalgComprehensiveSCFBufferizePass()";
+}
+
 def LinalgComprehensiveTensorBufferize :
     FunctionPass<"linalg-comprehensive-tensor-bufferize"> {
   let summary = "Bufferize (tensor into memref) tensor dialect ops.";
   let description = [{
     This pass bufferizes only tensor dialect ops.
     TODO: This will be moved to the bufferization dialect in the future.
   }];
   let constructor = "mlir::createLinalgComprehensiveTensorBufferizePass()";

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

 
 /// Specify that the value is known to bufferize to writable memory.
 void BufferizationAliasInfo::setBufferizesToWritableMemory(Value v) {
   bufferizeToWritableMemory.insert(v);
 }
 
 /// Return `true` if a value was marked as in-place bufferized.
 bool BufferizationAliasInfo::isInPlace(OpResult opResult) const {
-  bool inplace = inplaceBufferized.contains(opResult);
-#ifndef NDEBUG
-  if (inplace) {
-    auto bufferizableOp =
-        dyn_cast<BufferizableOpInterface>(opResult.getDefiningOp());
-    assert(bufferizableOp &&
-           "expected that in-place bufferized op is bufferizable");
-    SmallVector<OpOperand *> operands =
-        bufferizableOp.getAliasingOpOperand(opResult);
-    for (OpOperand *operand : operands)
-      assert(areAliasingBufferizedValues(operand->get(), opResult) &&
-             "expected that in-place bufferized OpResult aliases with "
-             "aliasing OpOperand");
-  }
-#endif // NDEBUG
-  return inplace;
+  return inplaceBufferized.contains(opResult);
 }
 
 /// Set the inPlace bufferization spec to true.
 void BufferizationAliasInfo::bufferizeInPlace(OpResult result,
                                               OpOperand &operand) {
   LLVM_DEBUG(llvm::dbgs() << "bufferizeInPlace: ");
   LLVM_DEBUG(result.print(llvm::dbgs()));
 
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(op->getContext());

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) 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 getResultBuffer(rewriter, forOp->getOpResult(index), state);
+          // return state.lookupBuffer(val);
+        });
+    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) {
+          Value toTensorOp =
+              rewriter.create<bufferization::ToTensorOp>(val.getLoc(), val);
+          state.mapBuffer(toTensorOp, val);
+          return toTensorOp;
+        });
+    iterArgs.insert(iterArgs.begin(), newForOp.getInductionVar());
+
+    // 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);
+
+    // Replace loop results.
+    // TODO: Some of this functionality should be moved to BufferizatioState for
+    // better code reuse.
+    for (OpOperand &use : forOp->getUses()) {
+      auto loopResult = use.get().cast<OpResult>();
+      int64_t resultIdx = loopResult.getResultNumber();
+      bool isTensorResult = indices.contains(resultIdx);
+      Value replacement = newForOp->getResult(resultIdx);
+      if (isTensorResult) {
+        rewriter.setInsertionPoint(use.getOwner());
+        replacement = rewriter.create<bufferization::ToTensorOp>(
+            loopResult.getLoc(), replacement);
+      }
+      use.set(replacement);
+      if (isTensorResult)
+        state.mapBuffer(replacement, newForOp->getResult(resultIdx));
+    }
 
-      // TODO: More general: Matching bbArg does not bufferize to a read.
-      Value resultBuffer = getResultBuffer(b, opResult, state);
-      if (!resultBuffer)
-        return failure();
+    // Erase the old loop.
+    forOp.erase();
+
+    // Bufferize loop body.
+    if (failed(comprehensive_bufferize::bufferize(loopBody, state)))
 
-      OpOperand &opOperand = forOp.getOpOperandForResult(opResult);
-      BlockArgument bbArg = forOp.getRegionIterArgForOpOperand(opOperand);
-      state.aliasInfo.createAliasInfoEntry(resultBuffer);
-      state.aliasInfo.insertNewBufferEquivalence(bbArg, resultBuffer);
-      state.mapBuffer(bbArg, resultBuffer);
-      state.mapBuffer(opResult, resultBuffer);
-    }
+    return success();
+};
 
-    // Bufferize loop body.
-    if (failed(comprehensive_bufferize::bufferize(&forOp.region(), state)))
-      return failure();
+LogicalResult mlir::linalg::comprehensive_bufferize::scf_ext::
+    AssertDestinationPassingStyle::run(Operation *op, BufferizationState &state,
+                                       SmallVector<Operation *> &newOps) {
+  LogicalResult status = success();
+  op->walk([&](scf::YieldOp yieldOp) {
+    auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
+    if (!forOp)
+      return WalkResult::advance();
 
-    // 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);
       if (!state.aliasInfo.areEquivalentBufferizedValues(operand.get(),
                                                          bbArg)) {
         // TODO: this could get resolved with copies but it can also turn into
         // swaps so we need to be careful about order of copies.
-        return yieldOp->emitError()
-               << "Yield operand #" << operand.getOperandNumber()
-               << " does not bufferize to an equivalent buffer to the matching"
-               << " enclosing scf::for operand";
+        status =
+            yieldOp->emitError()
+        return WalkResult::interrupt();
       }
-
-      // 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();
-  }
-};
+
+    return WalkResult::advance();
+  });
+  return status;
 
 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     scf::YieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
   }
 

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

options.allocationFns->memCpyFn = [](OpBuilder &b, Location loc, Value from,
       b.create<linalg::CopyOp>(loc, from, to);
     };
 
     // Enable InitTensorOp elimination.
     options.addPostAnalysisStep<
         linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
     // TODO: Find a way to enable these steps automatically.
     options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
+    options.addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
     options.allowReturnMemref = this->allowReturnMemref;
     options.allowUnknownOps = this->allowUnknownOps;
     options.analysisFuzzerSeed = this->analysisFuzzerSeed;
     options.testAnalysisOnly = this->testAnalysisOnly;
   }
 
   void runCleanupPipeline(StringRef name, Operation *op) {
struct LinalgComprehensiveFunctionBufferize
   LinalgComprehensiveFunctionBufferize() {}
 
   LinalgComprehensiveFunctionBufferize(
       const LinalgComprehensiveFunctionBufferize &p) {}
 
   void runOnFunction() final;
 };
 
+struct LinalgComprehensiveSCFBufferize
+    : public LinalgComprehensiveSCFBufferizeBase<
+          LinalgComprehensiveSCFBufferize> {
+  LinalgComprehensiveSCFBufferize() {}
+
+  LinalgComprehensiveSCFBufferize(const LinalgComprehensiveSCFBufferize &p) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
+                    scf::SCFDialect>();
+    bufferization_ext::registerBufferizableOpInterfaceExternalModels(registry);
+    scf_ext::registerBufferizableOpInterfaceExternalModels(registry);
+  }
+
+  void runOnFunction() override;
+};
+
 struct LinalgComprehensiveTensorBufferize
     : public LinalgComprehensiveTensorBufferizeBase<
           LinalgComprehensiveTensorBufferize> {
   LinalgComprehensiveTensorBufferize() {}
 
   LinalgComprehensiveTensorBufferize(
       const LinalgComprehensiveTensorBufferize &p) {}
 
void LinalgComprehensiveFunctionBufferize::runOnFunction() {
   }
 
   if (options.testAnalysisOnly)
     return;
 
   this->runCleanupPipeline("builtin.func", getOperation());
 }
 
+void LinalgComprehensiveSCFBufferize::runOnFunction() {
+  BufferizationOptions options;
+  options.allowReturnMemref = true;
+  options.allowUnknownOps = true;
+
+  if (failed(runComprehensiveBufferize(getOperation(), options)))
+    signalPassFailure();
+}
+
 void LinalgComprehensiveTensorBufferize::runOnFunction() {
   BufferizationOptions options;
   options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
   options.allowReturnMemref = true;
   options.allowUnknownOps = true;
 
   if (failed(runComprehensiveBufferize(getOperation(), options)))
     signalPassFailure();

 std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
   return std::make_unique<LinalgComprehensiveModuleBufferize>();
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveFunctionBufferizePass() {
   return std::make_unique<LinalgComprehensiveFunctionBufferize>();
 }
 
+std::unique_ptr<Pass> mlir::createLinalgComprehensiveSCFBufferizePass() {
+  return std::make_unique<LinalgComprehensiveSCFBufferize>();
+}
+
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveTensorBufferizePass() {
   return std::make_unique<LinalgComprehensiveTensorBufferize>();
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveVectorBufferizePass() {
   return std::make_unique<LinalgComprehensiveVectorBufferize>();
 }

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
 
 // Bufferize result IR of bufferization.
 // RUN: 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 -linalg-comprehensive-tensor-bufferize -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
+// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-scf-bufferize -canonicalize -split-input-file  | FileCheck %s --check-prefix=CHECK-SCF
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-vector-bufferize -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-VECTOR
 
+// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-scf-bufferize -canonicalize | \
+// RUN: mlir-opt -linalg-comprehensive-tensor-bufferize -canonicalize | FileCheck %s --check-prefix=CHECK-SCF-TENSOR
+
+// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-tensor-bufferize -canonicalize | \
+// RUN: mlir-opt -linalg-comprehensive-scf-bufferize -canonicalize | FileCheck %s --check-prefix=CHECK-SCF-TENSOR
+
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-tensor-bufferize -canonicalize | \
 // RUN: mlir-opt -linalg-comprehensive-vector-bufferize -canonicalize | FileCheck %s --check-prefix=CHECK-TENSOR-VECTOR
 
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-vector-bufferize -canonicalize | \
 // RUN: mlir-opt -linalg-comprehensive-tensor-bufferize -canonicalize | FileCheck %s --check-prefix=CHECK-VECTOR-TENSOR
 
 // CHECK-LABEL: func @use_of_unknown_op_1(
 //  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
func @writing_op_sequence(%t1: tensor<?xf32>, %s: f32,
   // CHECK-TENSOR-VECTOR: %[[write2_tensor:.*]] = bufferization.to_tensor %[[casted2]]
   // CHECK-TENSOR-VECTOR: return %[[write2_tensor]]
 
   // CHECK-VECTOR-TENSOR: %[[write2_tensor:.*]] = bufferization.to_tensor %[[casted2]]
   // CHECK-VECTOR-TENSOR: return %[[write2_tensor]]
   return %write2 : tensor<?xf32>
 }
 
+// -----
+
+// CHECK-SCF-LABEL: func @scf_for_equivalent_not_detected(
+//  CHECK-SCF-SAME:     %[[t1:.*]]: tensor<?xf32>
 
+// CHECK-SCF-TENSOR-LABEL: func @scf_for_equivalent_not_detected(
+//  CHECK-SCF-TENSOR-SAME:     %[[t1:.*]]: tensor<?xf32>
+func @scf_for_equivalent_not_detected(
+    %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]]) -> ({{.*}}) {
+
+  // CHECK-SCF-TENSOR: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
+  // CHECK-SCF-TENSOR: %[[alloc:.*]] = memref.alloc
+  // CHECK-SCF-TENSOR: %[[casted:.*]] = memref.cast %[[alloc]]
+  // CHECK-SCF-TENSOR: memref.copy %[[t1_memref]], %[[casted]]
+  // CHECK-SCF-TENSOR: %[[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]]
+
+    // CHECK-SCF-TENSOR: %[[alloc2:.*]] = memref.alloc
+    // CHECK-SCF-TENSOR: %[[casted2:.*]] = memref.cast %[[alloc2]]
+    // CHECK-SCF-TENSOR: memref.copy %[[arg0]], %[[casted2]]
+    // CHECK-SCF-TENSOR: memref.store %{{.*}}, %[[alloc2]]
+    %1 = tensor.insert %f into %arg0[%iv] : tensor<?xf32>
 
+    // CHECK-SCF: %[[insert_memref:.*]] = bufferization.to_memref %[[insert]]
+    // CHECK-SCF: scf.yield %[[insert_memref]]
+
+    // CHECK-SCF-TENSOR: scf.yield %[[casted2]]
+    scf.yield %1 : tensor<?xf32>
+  }
+  // CHECK-SCF: }
+  // CHECK-SCF-TENSOR: }
+
+  // CHECK-SCF: %[[scf_for_tensor:.*]] = bufferization.to_tensor %[[scf_for]]
+  // CHECK-SCF: return %[[scf_for_tensor]]
+
+  // CHECK-SCF-TENSOR: %[[scf_for_tensor:.*]] = bufferization.to_tensor %[[scf_for]]
+  // CHECK-SCF-TENSOR: 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.
 //===----------------------------------------------------------------------===//

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(