[mlir][bufferize] Extract buffer hoisting into separate function

[springerm]

Summary

This improves the modularity of the bufferization.

From now on, all ops that do not implement BufferizableOpInterface are considered hoisting barriers. Previously, all ops that do not implement the interface were not considered barriers and such ops had to be marked as barriers explicitly. This was unsafe because we could've hoisted across unknown ops where it was not safe to hoist.

As a side effect, this allows for cleaning up AffineBufferizableOpInterfaceImpl. This build unit no longer needed and can be deleted.

Depends On D121362

[springerm]

update

[springerm]

update

[gysit]

LGTM!

[springerm]

update

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;
 
   /// If set to `false`, all ops are bufferized (as long as they implement
   /// BufferizableOpInterface). Otherwise, only filtered ops are bufferized.
   bool hasFilter = false;
 
   /// A list of op filters that determine whether an op should be processed or

 /// allocated using `createAlloc`.
 LogicalResult createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
                             const BufferizationOptions &options);
 
 /// Creates a memcpy between two given buffers.
 LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from, Value to,
                            const BufferizationOptions &options);
 
-/// Finalize all buffer allocations, i.e., create alloc ops as specified in the
-/// bufferization options and deallocate all buffers.
+/// Finalize all buffer allocations.
+/// * Hoist buffer allocations as much as possible.
+/// * Create alloc/dealloc ops as specified by the bufferization options.
 LogicalResult finalizeBuffers(Operation *op,
                               const BufferizationOptions &options);
 } // namespace bufferization
 } // namespace mlir
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
-namespace mlir {
-namespace bufferization {
-
-/// AllocationHoistingBarrierOnly is an external implementation of
-/// BufferizableOpInterface for ops that are (not yet) bufferizable, but are
-/// known to be allocation hoisting barriers. All interface methods (except for
-/// `isAllocationHoistingBarrier`) are implemented conservatively.
-template <typename OpTy>
-struct AllocationHoistingBarrierOnly
-    : public BufferizableOpInterface::ExternalModel<
-          AllocationHoistingBarrierOnly<OpTy>, OpTy> {
-  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const AnalysisState &state) const {
-    return true;
-  }
-
-  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const AnalysisState &state) const {
-    return true;
-  }
-
-  SmallVector<OpOperand *>
-  getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const AnalysisState &state) const {
-    return {};
-  }
-
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
-    return {};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::None;
-  }
-
-  bool isWritable(Operation *op, Value value,
-                  const AnalysisState &state) const {
-    return false;
-  }
-
-  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
-    return failure();
-  }
-
-  bool isAllocationHoistingBarrier(Operation *op) const { return true; }
-};
-
-} // namespace bufferization
-} // namespace mlir
-
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

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

-//===- AffineInterfaceImpl.h - Affine 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
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_AFFINEINTERFACEIMPL_H
-#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_AFFINEINTERFACEIMPL_H
-
-namespace mlir {
-
-class DialectRegistry;
-
-namespace linalg {
-namespace comprehensive_bufferize {
-namespace affine_ext {
-
-void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
-
-} // namespace affine_ext
-} // namespace comprehensive_bufferize
-} // namespace linalg
-} // namespace mlir
-
-#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_AFFINEINTERFACEIMPL_H

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

bufferization::createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
   if (options.deallocationFn)
     return (*options.deallocationFn)(b, loc, allocatedBuffer);
 
   // Default buffer deallocation via DeallocOp.
   b.create<memref::DeallocOp>(loc, allocatedBuffer);
   return success();
 }
 
-/// Move the insertion point of the given builder to the beginning of a
-/// surrounding block as much as possible, while not crossing any allocation
-/// hoisting barriers.
-static void moveInsertionPointToAllocationHoistingBarrier(OpBuilder &b) {
-  Operation *op = b.getInsertionBlock()->getParentOp();
-  while (op) {
-    if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
-      if (bufferizableOp.isAllocationHoistingBarrier())
-        break;
-    op = op->getParentOp();
-  }
-
-  if (!op) {
-    // No allocation hoisting barrier found. Hoist to FuncOp.
-    op = b.getInsertionBlock()->getParentOp();
-    if (!isa<FuncOp>(op))
-      op = op->getParentOfType<FuncOp>();
-    assert(op && "could not find enclosing FuncOp");
-  }
-
-  // TODO: Handle cases where allocation hoisting barrier has more than one
-  // region or block.
-  assert(op->getNumRegions() == 1 &&
-         "allocation hoisting barriers with >1 regions not supported");
-  assert(op->getRegion(0).getBlocks().size() == 1 &&
-         "allocation hoisting barriers with >1 blocks not supported");
-  b.setInsertionPointToStart(&(op->getRegion(0).front()));
-}
-
 /// Compute the type of the `memref` to use for allocating the buffer for
 /// `shapedValue`. Also returns (by reference in `dynShape`), the value for the
-/// dynamic dimensions in the returned `memref` type. The function may also set
-/// the insertion point to an earlier location, where the allocation should
-/// happen ("allocation hoisting").
+/// dynamic dimensions in the returned `memref` type.
 static MemRefType getAllocationTypeAndShape(OpBuilder &b, Location loc,
                                             Value shapedValue,
                                             SmallVectorImpl<Value> &dynShape) {
   MemRefType allocMemRefType =
       getContiguousMemRefType(shapedValue.getType().cast<ShapedType>());
 
   // Compute the dynamic part of the shape.
   bool reifiedShapes = false;
for (const auto &dim : enumerate(allocMemRefType.getShape()))
         assert((shapedValue.getType().isa<UnrankedMemRefType>() ||
                 shapedValue.getType().isa<MemRefType>()) &&
                "expected MemRef type");
         dynShape.push_back(
             b.create<memref::DimOp>(loc, shapedValue, dim.index()));
       }
   }
 
-  // If the buffer is statically shaped, try to hoist it to the first enclosing
-  // parallel region.
-  // TODO: also hoist in the dynamic case. For now this relies on subsequent
-  // calls to LICM and buffer hoisting which will most likely not succeed.
-  // TODO: when packing, allocate a static bounding box which will enable more
-  // hoisting.
-  if (dynShape.empty())
-    moveInsertionPointToAllocationHoistingBarrier(b);
-
   return allocMemRefType;
 }
 
 static Value createBufferAllocation(OpBuilder &b, Location loc, MemRefType type,
                                     ValueRange dynShape) {
   auto allocaOp = b.create<memref::AllocaOp>(loc, type, dynShape);
   allocaOp->setAttr(kBufferAllocationAttr, b.getUnitAttr());
   return allocaOp.getResult();
 }
 
 /// Create an allocation after `shapedValue.getDefiningOp` (or at the top of the
 /// block in case of a bbArg).
 FailureOr<Value> BufferizationState::createAlloc(OpBuilder &b, Location loc,
                                                  Value shapedValue) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
   assert(shapedValue.getType().isa<ShapedType>());
   MemRefType memRefType = shapedValue.getType().dyn_cast<MemRefType>();
   SmallVector<Value> dynShape;
-  // Note: getAllocationTypeAndShape also sets the insertion point.
   MemRefType allocMemRefType =
       getAllocationTypeAndShape(b, loc, shapedValue, dynShape);
   Value alloc = createBufferAllocation(b, loc, allocMemRefType, dynShape);
   if (memRefType && memRefType != allocMemRefType) {
     assert(memref::CastOp::areCastCompatible(alloc.getType(), memRefType) &&
            "createAlloc: cast incompatible");
     alloc = b.create<memref::CastOp>(loc, memRefType, alloc);
   }
LogicalResult bufferization::createMemCpy(OpBuilder &b, Location loc,
                                           const BufferizationOptions &options) {
   if (options.memCpyFn)
     return (*options.memCpyFn)(b, loc, from, to);
 
   b.create<memref::CopyOp>(loc, from, to);
   return success();
 }
 
-LogicalResult
-bufferization::finalizeBuffers(Operation *op,
-                               const BufferizationOptions &options) {
+static LogicalResult
+createAllocDeallocOps(Operation *op, const BufferizationOptions &options) {
   IRRewriter rewriter(op->getContext());
 
   // Bufferization creates memref.alloca ops. After bufferization, these must be
   // rewritten to alloc/dealloc ops as specified in the bufferization options.
   WalkResult status = op->walk([&](memref::AllocaOp allocaOp) {
     // Ignore memref.alloca ops that were not created by the bufferization.
     if (!allocaOp->hasAttr(kBufferAllocationAttr))
       return WalkResult::skip();
if (failed(createDealloc(rewriter, alloc->getLoc(), *alloc, options)))
       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.
+static LogicalResult
+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();
+}
+
+LogicalResult
+bufferization::finalizeBuffers(Operation *op,
+                               const BufferizationOptions &options) {
+  if (failed(hoistBufferAllocations(op, options)))
+    return failure();
+  if (failed(createAllocDeallocOps(op, options)))
+    return failure();
+
+  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/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.cpp

-//===- AffineInterfaceImpl.cpp - Affine 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/AffineInterfaceImpl.h"
-
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
-
-using namespace mlir::bufferization;
-
-void mlir::linalg::comprehensive_bufferize::affine_ext::
-    registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
-  // AffineParallelOp bufferization not implemented yet. However, never hoist
-  // memref allocations across AffineParallelOp boundaries.
-  registry.addOpInterface<AffineParallelOp,
-                          AllocationHoistingBarrierOnly<AffineParallelOp>>();
-}

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/CMakeLists.txt

-set(LLVM_OPTIONAL_SOURCES
-  AffineInterfaceImpl.cpp
-  ModuleBufferization.cpp
-)
-
-add_mlir_dialect_library(MLIRAffineBufferizableOpInterfaceImpl
-  AffineInterfaceImpl.cpp
-
-  LINK_LIBS PUBLIC
-  MLIRAffine
-  MLIRBufferization
-)
-
 add_mlir_dialect_library(MLIRModuleBufferization
   ModuleBufferization.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferization
   MLIRBufferizationTransforms
   MLIRFunc
   MLIRFuncTransforms
   MLIRIR
   MLIRMemRef
 )

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
   MLIRLinalgPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
-  MLIRAffineBufferizableOpInterfaceImpl
   MLIRAffineUtils
   MLIRAnalysis
   MLIRArithmetic
   MLIRArithmeticTransforms
   MLIRBufferization
   MLIRComplex
   MLIRFunc
   MLIRFuncToLLVM

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

 //===- ComprehensiveBufferize.cpp - Single pass bufferization -------------===//
 //
 // 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 "PassDetail.h"
 
 #include "mlir/Dialect/Arithmetic/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
struct LinalgComprehensiveModuleBufferize
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
                 memref::MemRefDialect, tensor::TensorDialect,
                 vector::VectorDialect, scf::SCFDialect,
                 arith::ArithmeticDialect, func::FuncDialect, AffineDialect>();
-    affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
     arith::registerBufferizableOpInterfaceExternalModels(registry);
     linalg::registerBufferizableOpInterfaceExternalModels(registry);
     scf::registerBufferizableOpInterfaceExternalModels(registry);
     std_ext::registerModuleBufferizationExternalModels(registry);
     tensor::registerBufferizableOpInterfaceExternalModels(registry);
     vector::registerBufferizableOpInterfaceExternalModels(registry);
   }
 

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

 } // namespace mlir
 
 void mlir::scf::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addOpInterface<ExecuteRegionOp, ExecuteRegionOpInterface>();
   registry.addOpInterface<ForOp, ForOpInterface>();
   registry.addOpInterface<IfOp, IfOpInterface>();
   registry.addOpInterface<YieldOp, YieldOpInterface>();
-  registry
-      .addOpInterface<ParallelOp, AllocationHoistingBarrierOnly<ParallelOp>>();
 }

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/one-shot-bufferize-partial.mlir

 // CHECK-LABEL: func @unknown_op_may_read(
 func @unknown_op_may_read(%v: vector<5xf32>)
     -> (tensor<10xf32>, tensor<10xf32>) {
   %idx = arith.constant 0 : index
   %cst = arith.constant 5.0 : f32
 
   // One alloc for the init_tensor, another one because the transfer_write
   // bufferizes out-of-place.
-  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10xf32>
   // CHECK: %[[m1:.*]] = memref.alloc() {{.*}} : memref<10xf32>
+  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10xf32>
   %t1 = linalg.init_tensor [10] : tensor<10xf32>
 
   // CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[m1]]
   // CHECK: %[[filled_tensor:.*]] = bufferization.to_tensor %[[m1]]
   %filled = linalg.fill ins(%cst : f32) outs(%t1 : tensor<10xf32>) -> tensor<10xf32>
 
   // The transfer_write is out-of-place because "dummy_op" may read.
   // CHECK: memref.copy %[[m1]], %[[alloc]]

mlir/test/Dialect/Linalg/bufferize.mlir

#trait = {
   iterator_types = ["parallel", "parallel", "reduction"]
 }
 
 // Check the bufferization of init tensors.
 
 // CHECK-LABEL:   func @generic_with_init_tensor(
 // CHECK-SAME:                                   %[[ARG0_TENSOR:.*]]: tensor<2x3x4xvector<3x4xi4>>,
 // CHECK-SAME:                                   %[[ARG1_TENSOR:.*]]: tensor<3x2xf32>) -> tensor<3x2xf32> {
-// CHECK:           %[[INIT_BUFFER:.*]] = memref.alloc() {{.*}} : memref<3x2xf32>
-// CHECK-DAG:           %[[ARG0_MEMREF:.*]] = bufferization.to_memref %[[ARG0_TENSOR]] : memref<2x3x4xvector<3x4xi4>>
-// CHECK-DAG:           %[[ARG1_MEMREF:.*]] = bufferization.to_memref %[[ARG1_TENSOR]] : memref<3x2xf32>
+// CHECK-DAG:       %[[INIT_BUFFER:.*]] = memref.alloc() {{.*}} : memref<3x2xf32>
+// CHECK-DAG:       %[[ARG0_MEMREF:.*]] = bufferization.to_memref %[[ARG0_TENSOR]] : memref<2x3x4xvector<3x4xi4>>
+// CHECK-DAG:       %[[ARG1_MEMREF:.*]] = bufferization.to_memref %[[ARG1_TENSOR]] : memref<3x2xf32>
 // CHECK:           memref.copy %[[ARG1_MEMREF]], %[[INIT_BUFFER]] : memref<3x2xf32> to memref<3x2xf32>
 // CHECK:           linalg.generic
 // CHECK-SAME:      ins(%[[ARG0_MEMREF]] : memref<2x3x4xvector<3x4xi4>>)
 // CHECK-SAME:      outs(%[[INIT_BUFFER]] : memref<3x2xf32>) {
 func @generic_with_init_tensor(%arg0: tensor<2x3x4xvector<3x4xi4>>,
   %arg1: tensor<3x2xf32>) -> (tensor<3x2xf32>) {
 
   %0 = linalg.generic #trait

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

 func @main() {
   //  CHECK-DAG:   %[[C0:.*]] = arith.constant 0{{.*}} : f32
   //  CHECK-DAG:   %[[C1:.*]] = arith.constant 1{{.*}} : f32
   //  CHECK-DAG:   %[[C2:.*]] = arith.constant 2{{.*}} : f32
   %v0 = arith.constant 0.0 : f32
   %v1 = arith.constant 1.0 : f32
   %v2 = arith.constant 2.0 : f32
 
-  // CHECK-NEXT:   %[[C:.*]] = memref.alloca() {alignment = 128 : i64} : memref<f32>
-  // CHECK-NEXT:   %[[B:.*]] = memref.alloca() {alignment = 128 : i64} : memref<64xf32>
   // CHECK-NEXT:   %[[A:.*]] = memref.alloca() {alignment = 128 : i64} : memref<64xf32>
+  // CHECK-NEXT:   %[[B:.*]] = memref.alloca() {alignment = 128 : i64} : memref<64xf32>
+  // CHECK-NEXT:   %[[C:.*]] = memref.alloca() {alignment = 128 : i64} : memref<f32>
   //  CHECK-DAG:   %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
   //  CHECK-DAG:   %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
   //  CHECK-DAG:   %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, #[[$DYN_0D_MAP]]>
   %A = linalg.init_tensor [64] : tensor<64xf32>
   %B = linalg.init_tensor [64] : tensor<64xf32>
   %C = linalg.init_tensor [] : tensor<f32>
 
   //  CHECK-DAG:   linalg.fill ins(%[[C1]] : f32) outs(%[[A]] : memref<64xf32>)

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

 func @main() {
   //  CHECK-DAG:   %[[C0:.*]] = arith.constant 0{{.*}} : f32
   //  CHECK-DAG:   %[[C1:.*]] = arith.constant 1{{.*}} : f32
   //  CHECK-DAG:   %[[C2:.*]] = arith.constant 2{{.*}} : f32
   %v0 = arith.constant 0.0 : f32
   %v1 = arith.constant 1.0 : f32
   %v2 = arith.constant 2.0 : f32
 
-  // CHECK-NEXT:   %[[C:.*]] = memref.alloc() {alignment = 128 : i64} : memref<f32>
-  // CHECK-NEXT:   %[[B:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
   // CHECK-NEXT:   %[[A:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
+  // CHECK-NEXT:   %[[B:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
+  // CHECK-NEXT:   %[[C:.*]] = memref.alloc() {alignment = 128 : i64} : memref<f32>
   //  CHECK-DAG:   %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
   //  CHECK-DAG:   %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
   //  CHECK-DAG:   %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, #[[$DYN_0D_MAP]]>
   %A = linalg.init_tensor [64] : tensor<64xf32>
   %B = linalg.init_tensor [64] : tensor<64xf32>
   %C = linalg.init_tensor [] : tensor<f32>
 
   //  CHECK-DAG:   linalg.fill ins(%[[C1]] : f32) outs(%[[A]] : memref<64xf32>)

mlir/test/Dialect/Tensor/bufferize.mlir

 func @tensor.from_elements_0d(%arg0: index) -> tensor<index> {
   %0 = tensor.from_elements %arg0 : tensor<index>
   return %0 : tensor<index>
 }
 
 // CHECK-LABEL:   func @tensor.from_elements_1d(
 // CHECK-SAME:                               %[[ELEM0:.*]]: index,
 // CHECK-SAME:                               %[[ELEM1:.*]]: index) -> tensor<2xindex> {
-// CHECK:           %[[C0:.*]] = arith.constant 0 : index
-// CHECK:           %[[C1:.*]] = arith.constant 1 : index
-// CHECK:           %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<2xindex>
+// CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<2xindex>
 // CHECK:           store %[[ELEM0]], %[[MEMREF]][%[[C0]]]
 // CHECK:           store %[[ELEM1]], %[[MEMREF]][%[[C1]]]
 // CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
 // CHECK:           return %[[RET]] : tensor<2xindex>
 func @tensor.from_elements_1d(%arg0: index, %arg1: index) -> tensor<2xindex> {
   %0 = tensor.from_elements %arg0, %arg1 : tensor<2xindex>
   return %0 : tensor<2xindex>
 }
 
 // CHECK-LABEL: func @tensor.from_elements_2d(
 // CHECK-SAME:      %[[ELEM0:.*]]: index, %[[ELEM1:.*]]: index)
 // CHECK-SAME:      -> tensor<3x2xindex> {
 // CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index
-// CHECK:         %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<3x2xindex>
+// CHECK-DAG:     %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<3x2xindex>
 // CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C0]], %[[C0]]]
 // CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C0]], %[[C1]]]
 // CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C1]], %[[C0]]]
 // CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C1]], %[[C1]]]
 // CHECK:         store %[[ELEM0]], %[[MEMREF]][%[[C2]], %[[C0]]]
 // CHECK:         store %[[ELEM1]], %[[MEMREF]][%[[C2]], %[[C1]]]
 // CHECK:         %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]]
 // CHECK:         return %[[RET]] : tensor<3x2xindex>

 // CHECK-DAG: %[[F9:.*]] = arith.constant 9.0
 // CHECK-DAG: %[[F10:.*]] = arith.constant 1.0{{0+}}e+01
 // CHECK-DAG: %[[F11:.*]] = arith.constant 1.1{{0+}}e+01
 
 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
 
-// CHECK: %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<3x2x2xf32>
+// CHECK-DAG: %[[MEMREF:.*]] = memref.alloc() {{.*}} : memref<3x2x2xf32>
 
 // CHECK: store %[[F0]], %[[MEMREF]][%[[C0]], %[[C0]], %[[C0]]]
 // CHECK: store %[[F1]], %[[MEMREF]][%[[C0]], %[[C0]], %[[C1]]]
 // CHECK: store %[[F2]], %[[MEMREF]][%[[C0]], %[[C1]], %[[C0]]]
 // CHECK: store %[[F3]], %[[MEMREF]][%[[C0]], %[[C1]], %[[C1]]]
 // CHECK: store %[[F4]], %[[MEMREF]][%[[C1]], %[[C0]], %[[C0]]]
 // CHECK: store %[[F5]], %[[MEMREF]][%[[C1]], %[[C0]], %[[C1]]]
 // CHECK: store %[[F6]], %[[MEMREF]][%[[C1]], %[[C1]], %[[C0]]]
func @tensor.insert_slice(%t1: tensor<?x?xf32>, %t2: tensor<?x10xf32>,
   //     CHECK: return %[[r]]
   return %0 : tensor<?x?xf32>
 }
 
 // CHECK-LABEL: func @tensor.insert(
 //  CHECK-SAME:     %[[t1:.*]]: tensor<5xf32>, %[[idx1:.*]]: index,
 //  CHECK-SAME:     %[[f:.*]]: f32
 func @tensor.insert(%t1: tensor<5xf32>, %idx1: index, %f: f32) -> tensor<5xf32> {
-  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
-  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<5xf32>
+  // CHECK-DAG: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
+  // CHECK-DAG: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<5xf32>
   // CHECK: memref.copy %[[m1]], %[[alloc]]
   // CHECK: memref.store %[[f]], %[[alloc]][%[[idx1]]]
   %0 = tensor.insert %f into %t1[%idx1] : tensor<5xf32>
 
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
   // CHECK: return %[[r]]
   return %0 : tensor<5xf32>
 }

mlir/test/lib/Dialect/Linalg/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRLinalgTestPasses
   TestLinalgCodegenStrategy.cpp
   TestLinalgElementwiseFusion.cpp
   TestLinalgFusionTransforms.cpp
   TestLinalgHoisting.cpp
   TestLinalgTransforms.cpp
   TestPadFusion.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRAffine
-  MLIRAffineBufferizableOpInterfaceImpl
   MLIRArithmetic
   MLIRArithmeticTransforms
   MLIRBufferization
   MLIRBufferizationTransforms
   MLIRGPUTransforms
   MLIRLinalg
   MLIRLinalgTransforms
   MLIRLLVMToLLVMIRTranslation

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

gentbl_cc_library(
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td",
     deps = [
         ":BufferizableOpInterfaceTdFiles",
     ],
 )
 
-cc_library(
-    name = "AffineBufferizableOpInterfaceImpl",
-    srcs = [
-        "lib/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.cpp",
-    ],
-    hdrs = [
-        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h",
-    ],
-    includes = ["include"],
-    deps = [
-        ":Affine",
-        ":BufferizationDialect",
-        "//llvm:Support",
-    ],
-)
-
 td_library(
     name = "LinalgDocTdFiles",
     srcs = ["include/mlir/Dialect/Linalg/IR/LinalgDoc.td"],
     includes = ["include"],
     deps = [
         ":LinalgOpsTdFiles",
         ":LinalgStructuredOpsTdFiles",
     ],
hdrs = [
         "include/mlir/Dialect/Linalg/Transforms/Hoisting.h",
         "include/mlir/Dialect/Linalg/Transforms/Transforms.h",
         "include/mlir/Dialect/Linalg/Utils/Utils.h",
     ],
     includes = ["include"],
     deps = [
         ":Affine",
         ":AffineAnalysis",
-        ":AffineBufferizableOpInterfaceImpl",
         ":AffineUtils",
         ":Analysis",
         ":ArithmeticDialect",
         ":ArithmeticTransforms",
         ":ArithmeticUtils",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":ComplexDialect",

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

 cc_library(
     name = "TestLinalg",
     srcs = glob(["lib/Dialect/Linalg/*.cpp"]),
     defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],
     includes = ["lib/Dialect/Test"],
     deps = [
         "//llvm:Support",
         "//mlir:Affine",
-        "//mlir:AffineBufferizableOpInterfaceImpl",
         "//mlir:ArithmeticDialect",
         "//mlir:ArithmeticTransforms",
         "//mlir:BufferizationDialect",
         "//mlir:BufferizationTransforms",
         "//mlir:FuncDialect",
         "//mlir:GPUDialect",
         "//mlir:IR",
         "//mlir:LinalgOps",