[mlir][bufferize] Remove hoisting functionality from One-Shot Bufferize

[springerm]

Summary

The same functionality is already provided by -buffer-hoisting and -buffer-loop-hoisting.

The purpose of this change is to remove finalizeBuffers, which made it difficult to extend the bufferization to custom buffer types.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

struct BufferizationOptions {
   ///
   /// If set to `false`, BufferizableOpInterface implementations can try to be
   /// smart and choose to alias with "in" operands or other operands. E.g., the
   /// result of a `linalg.generic` op could bufferize in-place with an "in"
   /// OpOperand if the corresponding "out" operand is not used within the
   /// computation. Whether this pays off or not can be very input IR-specific.
   bool alwaysAliasingWithDest = true;
 
-  /// If set to `true`, try to hoist allocations out of blocks as much as
-  /// possible. An allocation is not hoisted across allocation hoisting barriers
-  /// as indicated by `BufferizableOpInterface::isAllocationHoistingBarrier`.
-  ///
-  /// Examples of allocation hoisting barriers are parallel loops or ops where
-  /// SSA values cannot be captured from the outside.
-  bool hoistAllocations = true;
-
   /// Buffer alignment for new memory allocations.
   unsigned int bufferAlignment = 128;
 
   /// A list of op filters that determine whether an op should be processed or
   /// ignored by the bufferization. If the filter has an ALLOW rule, only ops
   /// that are allowed and not denied are bufferized. If the filter does not
   /// have an ALLOW rule, only ops that are not denied are bufferized.
   SmallVector<OpFilterEntry> opFilter;
BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
                                                    Attribute memorySpace = {});
 
 /// Return a MemRef type with a static identity layout (i.e., no layout map). If
 /// the given tensor type is unranked, return an unranked MemRef type.
 BaseMemRefType
 getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
                                       Attribute memorySpace = {});
 
-/// Try to hoist all new buffer allocations until the next hoisting barrier.
-LogicalResult hoistBufferAllocations(Operation *op,
-                                     const BufferizationOptions &options);
-
 /// Create alloc/dealloc ops as specified in the bufferization options. If
 /// `onlyLeakingAlloc`, only those buffer allocations are processed for which no
 /// buffer deallocation can be created. `changed` is set to `true` if the IR was
 /// modified.
 LogicalResult createAllocDeallocOps(Operation *op,
                                     const BufferizationOptions &options,
                                     bool onlyLeakingAllocs = false,
                                     bool *changed = nullptr);
 
 } // namespace bufferization
 } // namespace mlir
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
                       "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return value.isa<OpResult>();
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
-          Return `true` if the op is an allocation hoisting barrier. Buffer
-          allocations will never be beyond such ops. E.g., ops with certain
-          parallel semantics may be allocation hoisting barriers. The majority
-          of ops, however, is not a barrier. Therefore, this method returns
-          `false` by default.
-        }],
-        /*retType=*/"bool",
-        /*methodName=*/"isAllocationHoistingBarrier",
-        /*args=*/(ins),
-        /*methodBody=*/"",
-        /*defaultImplementation=*/[{
-          return false;
-        }]
-      >,
-      InterfaceMethod<
-        /*desc=*/[{
           Return `true` if the `uRead` and `uWrite` do not constitute a RaW
           conflict. If they are conflicting or if it is unknown whether they are
           conflicting, return `false`. This method will never be called with
           OpOperands that do not have a tensor type. At least one of the two
           given OpOperands belongs to this operation.
 
           This method can be implemented to specify custom RaW analysis rules.
           If this method returns `true` the given OpOperands are not considered

mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h

 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
 /// Reuse an existing `BufferizationState`.
 ///
 /// Note: This function overload is useful for extending the bufferization.
 LogicalResult bufferizeOp(Operation *op,
                           BufferizationState &bufferizationState);
 
-/// Finalize all buffer allocations.
-/// * Hoist buffer allocations as much as possible.
-/// * Create alloc/dealloc ops as specified by the bufferization options.
+/// Finalize all buffer allocations: Create alloc/dealloc ops as specified by
+/// the bufferization options.
 LogicalResult finalizeBuffers(Operation *op,
                               const BufferizationOptions &options);
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

if (failed(options.createDealloc(rewriter, alloc->getLoc(), *alloc)))
       return WalkResult::interrupt();
 
     return WalkResult::advance();
   });
 
   return success(!status.wasInterrupted());
 }
 
-/// Try to hoist all new buffer allocations until the next hoisting barrier.
-// TODO: Consolidate this function with the existing buffer hoisting pass.
-LogicalResult
-bufferization::hoistBufferAllocations(Operation *op,
-                                      const BufferizationOptions &options) {
-  // Nothing to do if allocation hoisting is deactivated.
-  if (!options.hoistAllocations)
-    return success();
-
-  // Gather all buffer allocations that were created by the bufferization.
-  SmallVector<Operation *> allocaOps;
-  op->walk([&](memref::AllocaOp allocaOp) {
-    if (allocaOp->hasAttr(kBufferAllocationAttr))
-      allocaOps.push_back(allocaOp);
-  });
-
-  for (Operation *allocaOp : allocaOps) {
-    // TODO: Hoisting of allocs with dynamic shape not implemented.
-    if (!allocaOp->getOpOperands().empty())
-      continue;
-
-    Operation *op = allocaOp->getParentOp();
-    while (op) {
-      if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op)) {
-        if (bufferizableOp.isAllocationHoistingBarrier()) {
-          break;
-        }
-      } else {
-        // Op is not bufferizable: It may not be safe to hoist across this op.
-        break;
-      }
-      op = op->getParentOp();
-    }
-
-    // FuncOp is an allocation hoisting barrier, so this should never happen.
-    assert(op && "allocation hoisting barrier not found");
-
-    // Nothing to do if the insertion point is in the same block.
-    if (op == allocaOp->getParentOp())
-      continue;
-
-    // `op` may have multiple blocks. Make sure that we insert in the right one.
-    SmallVector<Block *> blocks;
-    for (Region &r : op->getRegions())
-      for (Block &b : r.getBlocks())
-        blocks.push_back(&b);
-    auto *insertionBlock = llvm::find_if(
-        blocks, [&](Block *b) { return b->findAncestorOpInBlock(*allocaOp); });
-    assert(insertionBlock != blocks.end() && "owning block not found");
-
-    // Move to the beginning of the block.
-    allocaOp->moveBefore(&(*insertionBlock)->front());
-  }
-
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 bool bufferization::isFunctionArgument(Value value) {
   auto bbArg = value.dyn_cast<BlockArgument>();
   if (!bbArg)
     return false;

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

static bool hasTensorSemantics(Operation *op) {
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   return hasTensorResult || hasTensorOperand;
 }
 
 LogicalResult
 bufferization::finalizeBuffers(Operation *op,
                                const BufferizationOptions &options) {
-  // Hoist buffers.
-  if (failed(hoistBufferAllocations(op, options)))
-    return failure();
-
   // Create allocation ops for "leaking buffers", i.e., buffer allocations that
   // escape block boundaries. If there are no leaking allocs, `hasLeakingAllocs`
   // is set to `false`.
   bool hasLeakingAllocs = false;
   if (failed(createAllocDeallocOps(op, options, /*onlyLeakingAllocs=*/true,
                                    &hasLeakingAllocs)))
     return failure();
 

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

bool isWritable(Operation *op, Value value,
     // currently used for testing only.
     if (BoolAttr writable = funcOp.getArgAttrOfType<BoolAttr>(
             bbArg.getArgNumber(), BufferizationDialect::kWritableAttrName))
       return writable.getValue();
 
     // All function arguments are writable by default.
     return true;
   }
-
-  bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 };
 
 } // namespace func_ext
 } // namespace bufferization
 } // namespace mlir
 
 void mlir::bufferization::func_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
     func::CallOp::attachInterface<func_ext::CallOpInterface>(*ctx);
     func::FuncOp::attachInterface<func_ext::FuncOpInterface>(*ctx);
     func::ReturnOp::attachInterface<func_ext::ReturnOpInterface>(*ctx);
   });
 }

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
-
-  bool isAllocationHoistingBarrier(Operation *op) const {
-    // Allocations should not be hoisted out of AssumingOps.
-    return true;
-  }
 };
 
 /// Bufferization of shape.assuming_yield. Bufferized as part of their enclosing
 /// ops, so this is for analysis only.
 struct AssumingYieldOpInterface
     : public BufferizableOpInterface::ExternalModel<AssumingYieldOpInterface,
                                                     shape::AssumingOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,

mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir

 // RUN: mlir-opt -buffer-hoisting -split-input-file %s | FileCheck %s
 
 // This file checks the behaviour of BufferHoisting pass for moving Alloc
 // operations to their correct positions.
 
 // Test Case:
 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp
 //   \   /
 //    bb3
 // BufferHoisting expected behavior: It should move the existing AllocOp to
 // the entry block.
 
 // CHECK-LABEL: func @condBranch
 func.func @condBranch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cf.cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   cf.br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cf.br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
 // CHECK-NEXT: cf.cond_br
 
 // -----
 
 // Test Case:
 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp
 //   \   /
 //    bb3
 // BufferHoisting expected behavior: It should not move the existing AllocOp
 // to any other block since the alloc has a dynamic dependency to block argument
 // %0 in bb2.
 
 // CHECK-LABEL: func @condBranchDynamicType
 func.func @condBranchDynamicType(
   %arg0: i1,
   %arg1: memref<?xf32>,
   %arg2: memref<?xf32>,
   %arg3: index) {
   cf.cond_br %arg0, ^bb1, ^bb2(%arg3: index)
 ^bb1:
   cf.br ^bb3(%arg1 : memref<?xf32>)
 ^bb2(%0: index):
   %1 = memref.alloc(%0) : memref<?xf32>
   test.buffer_based in(%arg1: memref<?xf32>) out(%1: memref<?xf32>)
   cf.br ^bb3(%1 : memref<?xf32>)
 ^bb3(%2: memref<?xf32>):
   test.copy(%2, %arg2) : (memref<?xf32>, memref<?xf32>)
   return
 }
 
 // CHECK-NEXT: cf.cond_br
 //      CHECK: ^bb2
 //      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc(%[[IDX]])
 // CHECK-NEXT: test.buffer_based
 
 // -----
 
 // Test Case:
 //      bb0
 //     /    \
 //   bb1    bb2 <- Initial position of AllocOp
 //    |     /  \
 //    |   bb3  bb4
 //    |     \  /
 //    \     bb5
 //     \    /
 //       bb6
 //        |
 //       bb7
 // BufferHoisting expected behavior: It should not move the existing AllocOp
 // to any other block since the alloc has a dynamic dependency to block argument
 // %0 in bb2.
 
 // CHECK-LABEL: func @condBranchDynamicTypeNested
 func.func @condBranchDynamicTypeNested(
   %arg0: i1,
   %arg1: memref<?xf32>,
   %arg2: memref<?xf32>,
   %arg3: index) {
   cf.cond_br %arg0, ^bb1, ^bb2(%arg3: index)
 ^bb1:
   cf.br ^bb6(%arg1 : memref<?xf32>)
 ^bb2(%0: index):
   %1 = memref.alloc(%0) : memref<?xf32>
   test.buffer_based in(%arg1: memref<?xf32>) out(%1: memref<?xf32>)
   cf.cond_br %arg0, ^bb3, ^bb4
 ^bb3:
   cf.br ^bb5(%1 : memref<?xf32>)
 ^bb4:
   cf.br ^bb5(%1 : memref<?xf32>)
 ^bb5(%2: memref<?xf32>):
   cf.br ^bb6(%2 : memref<?xf32>)
 ^bb6(%3: memref<?xf32>):
   cf.br ^bb7(%3 : memref<?xf32>)
 ^bb7(%4: memref<?xf32>):
   test.copy(%4, %arg2) : (memref<?xf32>, memref<?xf32>)
   return
 }
 
 // CHECK-NEXT: cf.cond_br
 //      CHECK: ^bb2
 //      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc(%[[IDX]])
 // CHECK-NEXT: test.buffer_based
 
 // -----
 
 // Test Case:
 //    bb0
 //   /   \
 //  |    bb1 <- Initial position of AllocOp
 //   \   /
 //    bb2
 // BufferHoisting expected behavior: It should move the existing AllocOp to
 // the entry block.
 
 // CHECK-LABEL: func @criticalEdge
 func.func @criticalEdge(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cf.cond_br %arg0, ^bb1, ^bb2(%arg1 : memref<2xf32>)
 ^bb1:
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cf.br ^bb2(%0 : memref<2xf32>)
 ^bb2(%1: memref<2xf32>):
   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
 // CHECK-NEXT: cf.cond_br
 
 // -----
 
 // Test Case:
 //    bb0 <- Initial position of the first AllocOp
 //   /   \
 //  bb1  bb2
 //   \   /
 //    bb3 <- Initial position of the second AllocOp
 // BufferHoisting expected behavior: It shouldn't move the AllocOps.
 
 // CHECK-LABEL: func @ifElse
 func.func @ifElse(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cf.cond_br %arg0,
     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
   cf.br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
   cf.br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
 ^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
   %7 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%7: memref<2xf32>) out(%7: memref<2xf32>)
   test.copy(%7, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: test.buffer_based
 //      CHECK: cf.br ^bb3
 //      CHECK: cf.br ^bb3
 // CHECK-NEXT: ^bb3
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc()
 // CHECK-NEXT: test.buffer_based
 //      CHECK: test.copy(%[[ALLOC1]]
 // CHECK-NEXT: return
 
 // -----
 
 // Test Case: No users for buffer in if-else CFG
 //    bb0 <- Initial position of AllocOp
 //   /   \
 //  bb1  bb2
 //   \   /
 //    bb3
 // BufferHoisting expected behavior: It shouldn't move the AllocOp.
 
 // CHECK-LABEL: func @ifElseNoUsers
 func.func @ifElseNoUsers(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cf.cond_br %arg0,
     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
   cf.br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
   cf.br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
 ^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
   test.copy(%arg1, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: test.buffer_based
 
 // -----
 
 // Test Case:
 //      bb0 <- Initial position of the first AllocOp
 //     /    \
 //   bb1    bb2
 //    |     /  \
 //    |   bb3  bb4
 //    \     \  /
 //     \     /
 //       bb5 <- Initial position of the second AllocOp
 // BufferHoisting expected behavior: AllocOps shouldn't be moved.
 
 // CHECK-LABEL: func @ifElseNested
 func.func @ifElseNested(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cf.cond_br %arg0,
     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
   cf.br ^bb5(%1, %2 : memref<2xf32>, memref<2xf32>)
 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
   cf.cond_br %arg0, ^bb3(%3 : memref<2xf32>), ^bb4(%4 : memref<2xf32>)
 ^bb3(%5: memref<2xf32>):
   cf.br ^bb5(%5, %3 : memref<2xf32>, memref<2xf32>)
 ^bb4(%6: memref<2xf32>):
   cf.br ^bb5(%3, %6 : memref<2xf32>, memref<2xf32>)
 ^bb5(%7: memref<2xf32>, %8: memref<2xf32>):
   %9 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%7: memref<2xf32>) out(%9: memref<2xf32>)
   test.copy(%9, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: test.buffer_based
 //      CHECK: cf.br ^bb5
 //      CHECK: cf.br ^bb5
 //      CHECK: cf.br ^bb5
 // CHECK-NEXT: ^bb5
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc()
 // CHECK-NEXT: test.buffer_based
 
 // -----
 
 // Test Case: Dead operations in a single block.
 // BufferHoisting expected behavior: It shouldn't move the AllocOps.
 
 // CHECK-LABEL: func @redundantOperations
 func.func @redundantOperations(%arg0: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
   %1 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%0: memref<2xf32>) out(%1: memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: test.buffer_based
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc()
 // CHECK-NEXT: test.buffer_based
 
 // -----
 
 // Test Case:
 //                                     bb0
 //                                    /   \
 // Initial pos of the 1st AllocOp -> bb1  bb2 <- Initial pos of the 2nd AllocOp
 //                                    \   /
 //                                     bb3
 // BufferHoisting expected behavior: Both AllocOps should be moved to the
 // entry block.
 
 // CHECK-LABEL: func @moving_alloc_and_inserting_missing_dealloc
 func.func @moving_alloc_and_inserting_missing_dealloc(
   %cond: i1,
     %arg0: memref<2xf32>,
     %arg1: memref<2xf32>) {
   cf.cond_br %cond, ^bb1, ^bb2
 ^bb1:
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
   cf.br ^exit(%0 : memref<2xf32>)
 ^bb2:
   %1 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg0: memref<2xf32>) out(%1: memref<2xf32>)
   cf.br ^exit(%1 : memref<2xf32>)
 ^exit(%arg2: memref<2xf32>):
   test.copy(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %{{.*}} = memref.alloc()
 // CHECK-NEXT: %{{.*}} = memref.alloc()
 // CHECK-NEXT: cf.cond_br
 
 // -----
 
 // Test Case: Invalid position of the DeallocOp. There is a user after
 // deallocation.
 //   bb0
 //  /   \
 // bb1  bb2 <- Initial position of AllocOp
 //  \   /
 //   bb3
 // BufferHoisting expected behavior: It should move the AllocOp to the entry
 // block.
 
 // CHECK-LABEL: func @moving_invalid_dealloc_op_complex
 func.func @moving_invalid_dealloc_op_complex(
   %cond: i1,
     %arg0: memref<2xf32>,
     %arg1: memref<2xf32>) {
   cf.cond_br %cond, ^bb1, ^bb2
 ^bb1:
   cf.br ^exit(%arg0 : memref<2xf32>)
 ^bb2:
   %1 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg0: memref<2xf32>) out(%1: memref<2xf32>)
   memref.dealloc %1 : memref<2xf32>
   cf.br ^exit(%1 : memref<2xf32>)
 ^exit(%arg2: memref<2xf32>):
   test.copy(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %{{.*}} = memref.alloc()
 // CHECK-NEXT: cf.cond_br
 
 // -----
 
 // Test Case: Nested regions - This test defines a BufferBasedOp inside the
 // region of a RegionBufferBasedOp.
 // BufferHoisting expected behavior: The AllocOp for the BufferBasedOp should
 // remain inside the region of the RegiobBufferBasedOp. The AllocOp of the
 // RegionBufferBasedOp should be moved to the entry block.
 
 // CHECK-LABEL: func @nested_regions_and_cond_branch
 func.func @nested_regions_and_cond_branch(
   %arg0: i1,
   %arg1: memref<2xf32>,
   %arg2: memref<2xf32>) {
   cf.cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   cf.br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = memref.alloc() : memref<2xf32>
   test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
   ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
     %1 = memref.alloc() : memref<2xf32>
     test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
     %tmp1 = math.exp %gen1_arg0 : f32
     test.region_yield %tmp1 : f32
   }
   cf.br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 // CHECK-NEXT:   %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT:   cf.cond_br
 //      CHECK:   test.region_buffer_based
 //      CHECK:     %[[ALLOC1:.*]] = memref.alloc()
 // CHECK-NEXT:     test.buffer_based
 
 // -----
 
 // Test Case: nested region control flow
 // The alloc position of %1 does not need to be changed and flows through
 // both if branches until it is finally returned.
 
 // CHECK-LABEL: func @nested_region_control_flow
 func.func @nested_region_control_flow(
   %arg0 : index,
   %arg1 : index) -> memref<?x?xf32> {
   %0 = arith.cmpi eq, %arg0, %arg1 : index
   %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
   %2 = scf.if %0 -> (memref<?x?xf32>) {
     scf.yield %1 : memref<?x?xf32>
   } else {
     %3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
     scf.yield %1 : memref<?x?xf32>
   }
   return %2 : memref<?x?xf32>
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
 // CHECK-NEXT: %{{.*}} = scf.if
 //      CHECK: else
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
 
 // -----
 
 // Test Case: nested region control flow with a nested buffer allocation in a
 // divergent branch.
 // The alloc positions of %1 does not need to be changed. %3 is moved upwards.
 
 // CHECK-LABEL: func @nested_region_control_flow_div
 func.func @nested_region_control_flow_div(
   %arg0 : index,
   %arg1 : index) -> memref<?x?xf32> {
   %0 = arith.cmpi eq, %arg0, %arg1 : index
   %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
   %2 = scf.if %0 -> (memref<?x?xf32>) {
     scf.yield %1 : memref<?x?xf32>
   } else {
     %3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
     scf.yield %3 : memref<?x?xf32>
   }
   return %2 : memref<?x?xf32>
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
 // CHECK-NEXT: %{{.*}} = scf.if
 
 // -----
 
 // Test Case: deeply nested region control flow with a nested buffer allocation
 // in a divergent branch.
 // The alloc position of %1 does not need to be changed. Allocs %4 and %5 are
 // moved upwards.
 
 // CHECK-LABEL: func @nested_region_control_flow_div_nested
 func.func @nested_region_control_flow_div_nested(
   %arg0 : index,
   %arg1 : index) -> memref<?x?xf32> {
   %0 = arith.cmpi eq, %arg0, %arg1 : index
   %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
   %2 = scf.if %0 -> (memref<?x?xf32>) {
     %3 = scf.if %0 -> (memref<?x?xf32>) {
       scf.yield %1 : memref<?x?xf32>
     } else {
       %4 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
       scf.yield %4 : memref<?x?xf32>
     }
     scf.yield %3 : memref<?x?xf32>
   } else {
     %5 = memref.alloc(%arg1, %arg1) : memref<?x?xf32>
     scf.yield %5 : memref<?x?xf32>
   }
   return %2 : memref<?x?xf32>
 }
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
 // CHECK-NEXT: %[[ALLOC2:.*]] = memref.alloc(%arg1, %arg1)
 // CHECK-NEXT: %{{.*}} = scf.if
 
 // -----
 
 // Test Case: deeply nested region control flow with a nested buffer allocation
 // that has dependency within a nested region should not be moved outside of
 // this region.
 
 // CHECK-LABEL: func @nested_region_control_flow_div_nested_dependencies
 func.func @nested_region_control_flow_div_nested_dependencies(
   %arg0: i32,
   %arg1: i1,
   %arg2: index) -> memref<?x?xf32> {
   %0 = scf.if %arg1 -> (memref<?x?xf32>) {
     %1 = arith.constant 1 : i32
     %2 = arith.addi %arg0, %1 : i32
     %3 = arith.index_cast %2 : i32 to index
     %4 = memref.alloc(%arg2, %3) : memref<?x?xf32>
     scf.yield %4 : memref<?x?xf32>
   } else {
     %1 = arith.constant 2 : i32
     %2 = arith.addi %arg0, %1 : i32
     %3 = arith.index_cast %2 : i32 to index
     %4 = memref.alloc(%arg2, %3) : memref<?x?xf32>
     scf.yield %4 : memref<?x?xf32>
   }
   return %0 : memref<?x?xf32>
 }
 
 //      CHECK: (%[[ARG0:.*]]: {{.*}}
 // CHECK-NEXT: %{{.*}} = scf.if
 // CHECK-NEXT: %{{.*}} = arith.constant
 // CHECK-NEXT: %{{.*}} = arith.addi
 // CHECK-NEXT: %[[FUNC:.*]] = arith.index_cast
 // CHECK-NEXT: alloc(%arg2, %[[FUNC]])
 // CHECK-NEXT: scf.yield
 // CHECK-NEXT: } else {
 // CHECK-NEXT: %{{.*}} = arith.constant
 // CHECK-NEXT: %{{.*}} = arith.addi
 // CHECK-NEXT: %[[FUNC:.*]] = arith.index_cast
 // CHECK-NEXT: alloc(%arg2, %[[FUNC]])
 
 // -----
 
 // Test Case: nested region control flow within a region interface.
 // The alloc positions of %0 does not need to be changed.
 
 // CHECK-LABEL: func @inner_region_control_flow
 func.func @inner_region_control_flow(%arg0 : index) -> memref<?x?xf32> {
   %0 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
   %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
     ^bb0(%arg1 : memref<?x?xf32>):
       test.region_if_yield %arg1 : memref<?x?xf32>
   } else {
     ^bb0(%arg1 : memref<?x?xf32>):
       test.region_if_yield %arg1 : memref<?x?xf32>
   } join {
     ^bb0(%arg1 : memref<?x?xf32>):
       test.region_if_yield %arg1 : memref<?x?xf32>
   }
   return %1 : memref<?x?xf32>
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
 // CHECK-NEXT: {{.*}} test.region_if
 
 // -----
 
 // Test Case: nested region control flow within a region interface including an
 // allocation in a divergent branch.
 // The alloc positions of %0 does not need to be changed. %2 is moved upwards.
 
 // CHECK-LABEL: func @inner_region_control_flow_div
 func.func @inner_region_control_flow_div(
   %arg0 : index,
   %arg1 : index) -> memref<?x?xf32> {
   %0 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
   %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
     ^bb0(%arg2 : memref<?x?xf32>):
       test.region_if_yield %arg2 : memref<?x?xf32>
   } else {
     ^bb0(%arg2 : memref<?x?xf32>):
       %2 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
       test.region_if_yield %2 : memref<?x?xf32>
   } join {
     ^bb0(%arg2 : memref<?x?xf32>):
       test.region_if_yield %arg2 : memref<?x?xf32>
   }
   return %1 : memref<?x?xf32>
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
 // CHECK-NEXT: {{.*}} test.region_if
 
 // -----
 
 // Test Case: Alloca operations shouldn't be moved.
 
 // CHECK-LABEL: func @condBranchAlloca
 func.func @condBranchAlloca(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cf.cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   cf.br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = memref.alloca() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cf.br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: cf.cond_br
 //      CHECK: ^bb2
 //      CHECK: ^bb2
 // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca()
 // CHECK-NEXT: test.buffer_based
 
 // -----
 
 // Test Case: Alloca operations shouldn't be moved. The alloc operation also
 // shouldn't be moved analogously to the ifElseNested test.
 
 // CHECK-LABEL: func @ifElseNestedAlloca
 func.func @ifElseNestedAlloca(
   %arg0: i1,
   %arg1: memref<2xf32>,
   %arg2: memref<2xf32>) {
   %0 = memref.alloca() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cf.cond_br %arg0,
     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
   cf.br ^bb5(%1, %2 : memref<2xf32>, memref<2xf32>)
 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
   cf.cond_br %arg0, ^bb3(%3 : memref<2xf32>), ^bb4(%4 : memref<2xf32>)
 ^bb3(%5: memref<2xf32>):
   cf.br ^bb5(%5, %3 : memref<2xf32>, memref<2xf32>)
 ^bb4(%6: memref<2xf32>):
   cf.br ^bb5(%3, %6 : memref<2xf32>, memref<2xf32>)
 ^bb5(%7: memref<2xf32>, %8: memref<2xf32>):
   %9 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%7: memref<2xf32>) out(%9: memref<2xf32>)
   test.copy(%9, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca()
 // CHECK-NEXT: test.buffer_based
 //      CHECK: ^bb5
 //      CHECK: ^bb5
 //      CHECK: ^bb5
 // CHECK-NEXT: ^bb5
 // CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
 // CHECK-NEXT: test.buffer_based
 
 // -----
 
 // Test Case: Alloca operations shouldn't be moved. The alloc operation should
 // be moved in the beginning analogous to the nestedRegionsAndCondBranch test.
 
 // CHECK-LABEL: func @nestedRegionsAndCondBranchAlloca
 func.func @nestedRegionsAndCondBranchAlloca(
   %arg0: i1,
   %arg1: memref<2xf32>,
   %arg2: memref<2xf32>) {
   cf.cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   cf.br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = memref.alloc() : memref<2xf32>
   test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
   ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
     %1 = memref.alloca() : memref<2xf32>
     test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
     %tmp1 = math.exp %gen1_arg0 : f32
     test.region_yield %tmp1 : f32
   }
   cf.br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 // CHECK-NEXT:   %[[ALLOC:.*]] = memref.alloc()
 // CHECK-NEXT:   cf.cond_br
 //      CHECK:   test.region_buffer_based
 //      CHECK:     %[[ALLOCA:.*]] = memref.alloca()
 // CHECK-NEXT:     test.buffer_based
 
 // -----
 
 // Test Case: structured control-flow loop using a nested alloc.
 // The alloc positions of %3 will be moved upwards.
 
 // CHECK-LABEL: func @loop_alloc
 func.func @loop_alloc(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = arith.cmpi eq, %i, %ub : index
     %3 = memref.alloc() : memref<2xf32>
     scf.yield %3 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: {{.*}} scf.for
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc()
 
 // -----
 
 // Test Case: structured control-flow loop with a nested if operation using
 // a deeply nested buffer allocation.
 // The allocation %4 is not moved upwards.
 
 // CHECK-LABEL: func @loop_nested_if_alloc
 func.func @loop_nested_if_alloc(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>) -> memref<2xf32> {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = arith.cmpi eq, %i, %ub : index
     %3 = scf.if %2 -> (memref<2xf32>) {
       %4 = memref.alloc() : memref<2xf32>
       scf.yield %4 : memref<2xf32>
     } else {
       scf.yield %0 : memref<2xf32>
     }
     scf.yield %3 : memref<2xf32>
   }
   return %1 : memref<2xf32>
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: {{.*}} scf.for
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc()
 
 // -----
 
 // Test Case: several nested structured control-flow loops with a deeply nested
 // buffer allocation inside an if operation.
 // Same behavior is an loop_nested_if_alloc: The allocs are not moved upwards.
 
 // CHECK-LABEL: func @loop_nested_alloc
 func.func @loop_nested_alloc(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = scf.for %i2 = %lb to %ub step %step
       iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
       %3 = scf.for %i3 = %lb to %ub step %step
         iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
         %4 = memref.alloc() : memref<2xf32>
         %5 = arith.cmpi eq, %i, %ub : index
         %6 = scf.if %5 -> (memref<2xf32>) {
           %7 = memref.alloc() : memref<2xf32>
           scf.yield %7 : memref<2xf32>
         } else {
           scf.yield %iterBuf3 : memref<2xf32>
         }
         scf.yield %6 : memref<2xf32>
       }
       scf.yield %3 : memref<2xf32>
     }
     scf.yield %2 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc()
 //      CHECK: %[[ALLOC2:.*]] = memref.alloc()
 
 // -----
 
 // CHECK-LABEL: func @loop_nested_alloc_dyn_dependency
 func.func @loop_nested_alloc_dyn_dependency(
   %lb: index,
   %ub: index,
   %step: index,
   %arg0: index,
   %buf: memref<?xf32>,
   %res: memref<?xf32>) {
   %0 = memref.alloc(%arg0) : memref<?xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<?xf32> {
     %2 = scf.for %i2 = %lb to %ub step %step
       iter_args(%iterBuf2 = %iterBuf) -> memref<?xf32> {
       %3 = scf.for %i3 = %lb to %ub step %step
         iter_args(%iterBuf3 = %iterBuf2) -> memref<?xf32> {
         %5 = arith.cmpi eq, %i, %ub : index
         %6 = scf.if %5 -> (memref<?xf32>) {
           %7 = memref.alloc(%i3) : memref<?xf32>
           scf.yield %7 : memref<?xf32>
         } else {
           scf.yield %iterBuf3 : memref<?xf32>
         }
         scf.yield %6 : memref<?xf32>
       }
       scf.yield %3 : memref<?xf32>
     }
     scf.yield %0 : memref<?xf32>
   }
   test.copy(%1, %res) : (memref<?xf32>, memref<?xf32>)
   return
 }
 
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc({{.*}})

mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir

 // RUN: mlir-opt -buffer-loop-hoisting -split-input-file %s | FileCheck %s
 
 // This file checks the behavior of BufferLoopHoisting pass for moving Alloc
 // operations in their correct positions.
 
 // Test Case:
 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp
 //   \   /
 //    bb3
 // BufferLoopHoisting expected behavior: It should not move the AllocOp.
 
 // CHECK-LABEL: func @condBranch
 func.func @condBranch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cf.cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   cf.br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cf.br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: cf.cond_br
 //      CHECK: %[[ALLOC:.*]] = memref.alloc()
 
 // -----
 
 // Test Case:
 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp
 //   \   /
 //    bb3
 // BufferLoopHoisting expected behavior: It should not move the existing AllocOp
 // to any other block since the alloc has a dynamic dependency to block argument
 // %0 in bb2.
 
 // CHECK-LABEL: func @condBranchDynamicType
 func.func @condBranchDynamicType(
   %arg0: i1,
   %arg1: memref<?xf32>,
   %arg2: memref<?xf32>,
   %arg3: index) {
   cf.cond_br %arg0, ^bb1, ^bb2(%arg3: index)
 ^bb1:
   cf.br ^bb3(%arg1 : memref<?xf32>)
 ^bb2(%0: index):
   %1 = memref.alloc(%0) : memref<?xf32>
   test.buffer_based in(%arg1: memref<?xf32>) out(%1: memref<?xf32>)
   cf.br ^bb3(%1 : memref<?xf32>)
 ^bb3(%2: memref<?xf32>):
   test.copy(%2, %arg2) : (memref<?xf32>, memref<?xf32>)
   return
 }
 
 // CHECK-NEXT: cf.cond_br
 //      CHECK: ^bb2
 //      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc(%[[IDX]])
 // CHECK-NEXT: test.buffer_based
 
 // -----
 
 // Test Case: Nested regions - This test defines a BufferBasedOp inside the
 // region of a RegionBufferBasedOp.
 // BufferLoopHoisting expected behavior: The AllocOp for the BufferBasedOp
 // should remain inside the region of the RegionBufferBasedOp. The AllocOp of
 // the RegionBufferBasedOp should not be moved during this pass.
 
 // CHECK-LABEL: func @nested_regions_and_cond_branch
 func.func @nested_regions_and_cond_branch(
   %arg0: i1,
   %arg1: memref<2xf32>,
   %arg2: memref<2xf32>) {
   cf.cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   cf.br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = memref.alloc() : memref<2xf32>
   test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
   ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
     %1 = memref.alloc() : memref<2xf32>
     test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
     %tmp1 = math.exp %gen1_arg0 : f32
     test.region_yield %tmp1 : f32
   }
   cf.br ^bb3(%0 : memref<2xf32>)
 ^bb3(%1: memref<2xf32>):
   test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
   return
 }
 // CHECK-NEXT:   cf.cond_br
 //      CHECK:   %[[ALLOC0:.*]] = memref.alloc()
 //      CHECK:   test.region_buffer_based
 //      CHECK:     %[[ALLOC1:.*]] = memref.alloc()
 // CHECK-NEXT:     test.buffer_based
 
 // -----
 
 // Test Case: nested region control flow
 // The alloc position of %1 does not need to be changed and flows through
 // both if branches until it is finally returned.
 
 // CHECK-LABEL: func @nested_region_control_flow
 func.func @nested_region_control_flow(
   %arg0 : index,
   %arg1 : index) -> memref<?x?xf32> {
   %0 = arith.cmpi eq, %arg0, %arg1 : index
   %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
   %2 = scf.if %0 -> (memref<?x?xf32>) {
     scf.yield %1 : memref<?x?xf32>
   } else {
     %3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
     scf.yield %1 : memref<?x?xf32>
   }
   return %2 : memref<?x?xf32>
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc(%arg0, %arg0)
 // CHECK-NEXT: %{{.*}} = scf.if
 //      CHECK: else
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%arg0, %arg1)
 
 // -----
 
 // Test Case: structured control-flow loop using a nested alloc.
 // The alloc positions of %3 should not be changed.
 
 // CHECK-LABEL: func @loop_alloc
 func.func @loop_alloc(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = arith.cmpi eq, %i, %ub : index
     %3 = memref.alloc() : memref<2xf32>
     scf.yield %3 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: {{.*}} scf.for
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc()
 
 // -----
 
 // Test Case: structured control-flow loop with a nested if operation using
 // a deeply nested buffer allocation.
 // The allocation %4 should not be moved upwards due to a back-edge dependency.
 
 // CHECK-LABEL: func @loop_nested_if_alloc
 func.func @loop_nested_if_alloc(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>) -> memref<2xf32> {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = arith.cmpi eq, %i, %ub : index
     %3 = scf.if %2 -> (memref<2xf32>) {
       %4 = memref.alloc() : memref<2xf32>
       scf.yield %4 : memref<2xf32>
     } else {
       scf.yield %0 : memref<2xf32>
     }
     scf.yield %3 : memref<2xf32>
   }
   return %1 : memref<2xf32>
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: {{.*}} scf.for
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc()
 
 // -----
 
 // Test Case: several nested structured control-flow loops with deeply nested
 // buffer allocations inside an if operation.
 // Behavior: The allocs %0, %4 and %9 are moved upwards, while %7 and %8 stay
 // in their positions.
 
 // CHECK-LABEL: func @loop_nested_alloc
 func.func @loop_nested_alloc(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = scf.for %i2 = %lb to %ub step %step
       iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
       %3 = scf.for %i3 = %lb to %ub step %step
         iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
         %4 = memref.alloc() : memref<2xf32>
         %5 = arith.cmpi eq, %i, %ub : index
         %6 = scf.if %5 -> (memref<2xf32>) {
           %7 = memref.alloc() : memref<2xf32>
           %8 = memref.alloc() : memref<2xf32>
           scf.yield %8 : memref<2xf32>
         } else {
           scf.yield %iterBuf3 : memref<2xf32>
         }
         %9 = memref.alloc() : memref<2xf32>
         scf.yield %6 : memref<2xf32>
       }
       scf.yield %3 : memref<2xf32>
     }
     scf.yield %2 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc()
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc()
 // CHECK-NEXT: %[[ALLOC2:.*]] = memref.alloc()
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 //      CHECK: {{.*}} = scf.if
 //      CHECK: %[[ALLOC3:.*]] = memref.alloc()
 //      CHECK: %[[ALLOC4:.*]] = memref.alloc()
 
 // -----
 
 // CHECK-LABEL: func @loop_nested_alloc_dyn_dependency
 func.func @loop_nested_alloc_dyn_dependency(
   %lb: index,
   %ub: index,
   %step: index,
   %arg0: index,
   %buf: memref<?xf32>,
   %res: memref<?xf32>) {
   %0 = memref.alloc(%arg0) : memref<?xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<?xf32> {
     %2 = scf.for %i2 = %lb to %ub step %step
       iter_args(%iterBuf2 = %iterBuf) -> memref<?xf32> {
       %3 = scf.for %i3 = %lb to %ub step %step
         iter_args(%iterBuf3 = %iterBuf2) -> memref<?xf32> {
         %4 = memref.alloc(%i3) : memref<?xf32>
         %5 = arith.cmpi eq, %i, %ub : index
         %6 = scf.if %5 -> (memref<?xf32>) {
           %7 = memref.alloc(%i3) : memref<?xf32>
           scf.yield %7 : memref<?xf32>
         } else {
           scf.yield %iterBuf3 : memref<?xf32>
         }
         %8 = memref.alloc(%i3) : memref<?xf32>
         scf.yield %6 : memref<?xf32>
       }
       scf.yield %3 : memref<?xf32>
     }
     scf.yield %0 : memref<?xf32>
   }
   test.copy(%1, %res) : (memref<?xf32>, memref<?xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: {{.*}} = scf.for
 //      CHECK: %[[ALLOC1:.*]] = memref.alloc({{.*}})
 //      CHECK: %[[ALLOC2:.*]] = memref.alloc({{.*}})
 
 // -----
 
 // CHECK-LABEL: func @hoist_one_loop
 func.func @hoist_one_loop(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
       %2 = memref.alloc() : memref<2xf32>
       scf.yield %0 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
 
 // -----
 
 // CHECK-LABEL: func @no_hoist_one_loop
 func.func @no_hoist_one_loop(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
       %1 = memref.alloc() : memref<2xf32>
       scf.yield %1 : memref<2xf32>
   }
   test.copy(%0, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: {{.*}} = scf.for
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 
 // -----
 
 // CHECK-LABEL: func @hoist_multiple_loop
 func.func @hoist_multiple_loop(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = scf.for %i2 = %lb to %ub step %step
       iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
         %3 = memref.alloc() : memref<2xf32>
         scf.yield %0 : memref<2xf32>
     }
     scf.yield %0 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
 
 // -----
 
 // CHECK-LABEL: func @no_hoist_one_loop_conditional
 func.func @no_hoist_one_loop_conditional(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
       %1 = arith.cmpi eq, %i, %ub : index
       %2 = scf.if %1 -> (memref<2xf32>) {
         %3 = memref.alloc() : memref<2xf32>
         scf.yield %3 : memref<2xf32>
       } else {
         scf.yield %iterBuf : memref<2xf32>
       }
     scf.yield %2 : memref<2xf32>
   }
   test.copy(%0, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: {{.*}} = scf.for
 //      CHECK: {{.*}} = scf.if
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 
 // -----
 
 // CHECK-LABEL: func @hoist_one_loop_conditional
 func.func @hoist_one_loop_conditional(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = arith.cmpi eq, %lb, %ub : index
   %2 = scf.if %1 -> (memref<2xf32>) {
     %3 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
       %4 = memref.alloc() : memref<2xf32>
       scf.yield %0 : memref<2xf32>
     }
     scf.yield %0 : memref<2xf32>
   }
   else
   {
     scf.yield %0 : memref<2xf32>
   }
   test.copy(%2, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: {{.*}} = scf.if
 // CHECK-NEXT: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 //      CHECK: {{.*}} = scf.for
 
 // -----
 
 // CHECK-LABEL: func @no_hoist_one_loop_dependency
 func.func @no_hoist_one_loop_dependency(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
       %2 = memref.alloc(%i) : memref<?xf32>
       scf.yield %0 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
 
 // -----
 
 // CHECK-LABEL: func @partial_hoist_multiple_loop_dependency
 func.func @partial_hoist_multiple_loop_dependency(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = scf.for %i2 = %lb to %ub step %step
       iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
         %3 = memref.alloc(%i) : memref<?xf32>
         scf.yield %0 : memref<2xf32>
     }
     scf.yield %0 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
 
 // -----
 
 // Test with allocas to ensure that op is also considered.
 
 // CHECK-LABEL: func @hoist_alloca
 func.func @hoist_alloca(
   %lb: index,
   %ub: index,
   %step: index,
   %buf: memref<2xf32>,
   %res: memref<2xf32>) {
   %0 = memref.alloca() : memref<2xf32>
   %1 = scf.for %i = %lb to %ub step %step
     iter_args(%iterBuf = %buf) -> memref<2xf32> {
     %2 = scf.for %i2 = %lb to %ub step %step
       iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
         %3 = memref.alloca() : memref<2xf32>
         scf.yield %0 : memref<2xf32>
     }
     scf.yield %0 : memref<2xf32>
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOCA0:.*]] = memref.alloca({{.*}})
 // CHECK-NEXT: %[[ALLOCA1:.*]] = memref.alloca({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for

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

-// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -buffer-loop-hoisting -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP

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

 
 // CHECK-LABEL: func @scf_while_non_equiv_condition(
 //  CHECK-SAME:     %[[arg0:.*]]: memref<5xi1, #{{.*}}>, %[[arg1:.*]]: memref<5xi1, #{{.*}}>
 func.func @scf_while_non_equiv_condition(%arg0: tensor<5xi1>,
                                          %arg1: tensor<5xi1>,
                                          %idx: index)
   -> (tensor<5xi1>, tensor<5xi1>)
 {
-  // These allocation used to be inside the scf.while loop, but they were
-  // hoisted.
-  // CHECK: %[[a0:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-  // CHECK: %[[a1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-  // CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[arg0]], %[[w1:.*]] = %[[arg1]]) {{.*}} {
+  // CHECK: %[[clone1:.*]] = bufferization.clone %[[arg1]]
+  // CHECK: %[[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: %[[casted1:.*]] = memref.cast %[[a1]]
+    // CHECK: memref.dealloc %[[w1]]
+    // CHECK: %[[a0:.*]] = memref.alloc() {{.*}} : memref<5xi1>
     // CHECK: memref.copy %[[w0]], %[[a0]]
+    // CHECK: memref.dealloc %[[w0]]
     // CHECK: %[[casted0:.*]] = memref.cast %[[a0]]
+    // CHECK: %[[casted1:.*]] = memref.cast %[[a1]]
     // CHECK: scf.condition(%[[condition]]) %[[casted1]], %[[casted0]]
     %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-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>)
 {
-  // These allocation used to be inside the scf.while loop, but they were
-  // hoisted.
-  // CHECK: %[[a0:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-  // CHECK: %[[a1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-  // CHECK: %[[a2:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-  // CHECK: %[[a3:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-  // CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[arg0]], %[[w1:.*]] = %[[arg1]]) {{.*}} {
+  // CHECK: %[[clone1:.*]] = bufferization.clone %[[arg1]]
+  // CHECK: %[[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: memref.copy %[[w1]], %[[a3]]
-    // CHECK: %[[casted3:.*]] = memref.cast %[[a3]]
-    // CHECK: memref.copy %[[w0]], %[[a2]]
-    // CHECK: %[[casted2:.*]] = memref.cast %[[a2]]
-    // CHECK: scf.condition(%[[condition]]) %[[casted3]], %[[casted2]]
+    // 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: %[[casted0:.*]] = memref.cast %[[a0]]
+    // CHECK: %[[casted1:.*]] = memref.cast %[[a1]]
+    // CHECK: scf.condition(%[[condition]]) %[[casted1]], %[[casted0]]
     %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: memref.copy %[[b1]], %[[a1]]
-    // CHECK: %[[casted1:.*]] = memref.cast %[[a1]]
-    // CHECK: memref.copy %[[b0]], %[[a0]]
-    // CHECK: %[[casted0:.*]] = memref.cast %[[a0]]
-    // CHECK: scf.yield %[[casted1]], %[[casted0]]
+    // CHECK: %[[a3:.*]] = memref.alloc() {{.*}} : memref<5xi1>
+    // CHECK: memref.copy %[[b1]], %[[a3]]
+    // CHECK: memref.dealloc %[[b1]]
+    // CHECK: %[[a2:.*]] = memref.alloc() {{.*}} : memref<5xi1>
+    // CHECK: memref.copy %[[b0]], %[[a2]]
+    // CHECK: %[[casted2:.*]] = memref.cast %[[a2]]
+    // CHECK: %[[casted3:.*]] = memref.cast %[[a3]]
+    // CHECK: scf.yield %[[casted3]], %[[casted2]]
     // 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-DAG: memref.dealloc %[[a0]]
-  // CHECK-DAG: memref.dealloc %[[a1]]
   // CHECK: return %[[loop]]#0, %[[loop]]#1
   return %r0, %r1 : tensor<5xi1>, tensor<5xi1>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_while_iter_arg_result_mismatch(
 //  CHECK-SAME:     %[[arg0:.*]]: memref<5xi1, #{{.*}}>, %[[arg1:.*]]: memref<5xi1, #{{.*}}>
-//       CHECK:   %[[alloc1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
 //       CHECK:   %[[alloc2:.*]] = memref.alloc() {{.*}} : memref<5xi1>
-//       CHECK:   scf.while (%[[arg3:.*]] = %[[arg1]]) : (memref<5xi1, #{{.*}}) -> () {
+//       CHECK:   %[[clone:.*]] = bufferization.clone %[[arg1]]
+//       CHECK:   scf.while (%[[arg3:.*]] = %[[clone]]) : (memref<5xi1, #{{.*}}) -> () {
+//       CHECK:     memref.dealloc %[[arg3]]
 //       CHECK:     %[[load:.*]] = memref.load %[[arg0]]
 //       CHECK:     scf.condition(%[[load]])
 //       CHECK:   } do {
 //       CHECK:     memref.copy %[[arg0]], %[[alloc2]]
 //       CHECK:     memref.store %{{.*}}, %[[alloc2]]
+//       CHECK:     %[[alloc1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
 //       CHECK:     memref.copy %[[alloc2]], %[[alloc1]]
 //       CHECK:     %[[casted:.*]] = memref.cast %[[alloc1]] : memref<5xi1> to memref<5xi1, #{{.*}}>
 //       CHECK:     scf.yield %[[casted]]
 //       CHECK:   }
-//   CHECK-DAG:   memref.dealloc %[[alloc1]]
 //   CHECK-DAG:   memref.dealloc %[[alloc2]]
 func.func @scf_while_iter_arg_result_mismatch(%arg0: tensor<5xi1>,
                                               %arg1: tensor<5xi1>,
                                               %arg2: index) {
   scf.while (%arg3 = %arg1) : (tensor<5xi1>) -> () {
     %0 = tensor.extract %arg0[%arg2] : tensor<5xi1>
     scf.condition(%0)
   } do {
     %0 = "dummy.some_op"() : () -> index
     %1 = "dummy.another_op"() : () -> i1
     %2 = tensor.insert %1 into %arg0[%0] : tensor<5xi1>
     scf.yield %2 : tensor<5xi1>
   }
   return
 }