[mlir][bufferize][NFC] Rename `comprehensive-function-bufferize` to `one-shot-bufferize`

[springerm]

Summary

The related functionality is moved over to the bufferization dialect. Test cases are cleaned up a bit.

Depends On D120190

[pifon2a]

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

168

did you intend to pass options by value?

[springerm]

rebase

[springerm]

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

168

yes i did

[jpienaar]

This seems to be causing build failures: https://buildkite.com/mlir/mlir-core/builds/20561#46e3ef4e-7afa-4b90-a9c3-27c112373d5e

[springerm]

Potential fix: D120326

[mehdi_amini]

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

209

This was maybe fine for a test pass, but I really don't think we should do this kind of things in a an actual pass.

For example, I'm currently chasing a bug that I narrowed down to this pass, and doing more A/B comparison it turns out that the output of the pass is identical in both case for my path, but this runPipeline differ. However I have no tool to bisect further!

passes should be small units that composes, a cleanup pipeline inside a pass is a anti-pattern.

[springerm]

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

209

D155923

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

private:
   BufferizationAliasInfo aliasInfo;
 };
 
 /// Analyze `op` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeOp(Operation *op, AnalysisBufferizationState &state);
 
 /// Run One-Shot Bufferize on the given op: Analysis + Bufferization
-LogicalResult
-runOneShotBufferize(Operation *op,
-                    std::unique_ptr<AnalysisBufferizationOptions> options);
+LogicalResult runOneShotBufferize(Operation *op,
+                                  const AnalysisBufferizationOptions &options);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H

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

 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace bufferization {
+struct AnalysisBufferizationOptions;
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 /// Creates an instance of the BufferDeallocation pass to free all allocated
 /// buffers.
 std::unique_ptr<Pass> createBufferDeallocationPass();
 
 /// Creates a pass that moves allocations upwards to reduce the number of
 /// required copies that are inserted during the BufferDeallocation pass.
 std::unique_ptr<Pass> createBufferHoistingPass();
 
 /// Creates a pass that moves allocations upwards out of loops. This avoids
 /// reallocations inside of loops.
 std::unique_ptr<Pass> createBufferLoopHoistingPass();
 
 /// Creates a pass that converts memref function results to out-params.
 std::unique_ptr<Pass> createBufferResultsToOutParamsPass();
 
 /// Creates a pass that finalizes a partial bufferization by removing remaining
 /// bufferization.to_tensor and bufferization.to_memref operations.
 std::unique_ptr<OperationPass<FuncOp>> createFinalizingBufferizePass();
 
+/// Create a pass that bufferizes all ops that implement BufferizableOpInterface
+/// with One-Shot Bufferize.
+std::unique_ptr<Pass> createOneShotBufferizePass();
+
+/// Create a pass that bufferizes all ops that implement BufferizableOpInterface
+/// with One-Shot Bufferize and the specified bufferization options.
+std::unique_ptr<Pass>
+createOneShotBufferizePass(const AnalysisBufferizationOptions &options);
+
 /// Creates a pass that promotes heap-based allocations to stack-based ones.
 /// Only buffers smaller than the provided size are promoted.
 /// Dynamic shaped buffers are promoted up to the given rank.
 std::unique_ptr<Pass>
 createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes = 1024,
                                 unsigned maxRankOfAllocatedMemRef = 1);
 
 /// Creates a pass that promotes heap-based allocations to stack-based ones.

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

let description = [{
     `bufferization.to_buffer` operations.
 
     This pass will fail if not all operations can be removed or if any operation
     with tensor typed operands remains.
   }];
   let constructor = "mlir::bufferization::createFinalizingBufferizePass()";
 }
 
+def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
+  let summary = "One-Shot Bufferize";
+  let description = [{
+    This pass bufferizes all ops that implement `BufferizableOpInterface`. It
+    first performs an inplacability analysis on SSA use-def chains of tensor
+    values to determine which OpOperands may bufferize in-place, i.e., without
+    inserting a buffer copy. It then rewrites the IR, inserting a buffer
+    allocation and copy for each OpOperand that was decided to bufferize
+    out-of-place.
+
+    One-Shot Bufferize (and `BufferizableOpInterface`) was designed for ops that
+    are in destination-passing style. When bufferizing such ops, it is possible
+    to reuse the buffer of a tensor OpOperand for a tensor OpResult. In essence,
+    a possible destination of an operation is already passed as an SSA value.
+
+    `tensor.insert` is an example for an op in destination-passing style. E.g.,
+    when bufferizing `%t0 = tensor.insert %f into %dest[%idx]`, `buffer(%t0)` is
+    identical to `buffer(%dest)` in the absence of RaW conflicts. As a counter
+    example, `tensor.generate` is not in destination-passing style and always
+    results in a new buffer allocation.
+
+    One-Shot Bufferize deallocates all buffers that it allocates. Yielding newly
+    allocated buffers from a block is not supported yet and such IR will be
+    rejected. For testing purposes and compatibility with partial bufferization,
+    One-Shot Bufferize can be run with `allow-return-memref=1 create-dealloc=0`
+    to allow such IR.
+
+    One-Shot Bufferize will by default reject IR that contains non-bufferizable
+    op, i.e., ops that do not implemement BufferizableOpInterface. Such IR can
+    be allowed with `allow-unknown-ops=1`. In that case, to_memref and to_tensor
+    ops will be generated at the bufferization boundary. This is useful for
+    compatibility with existing partial bufferization passes: These can
+    bufferize the remaining IR after running One-Shot Bufferize.
+
+    Note: Running One-Shot Bufferize after a partial bufferization pass is
+    currently not supported. Running partial bufferization passes after running
+    One-Shot Bufferize is supported and the recommended way to gradually
+    migrate from partial bufferization to One-Shot Bufferize.
+
+    With `dialect-filter`, bufferization can be restricted to a set of dialects.
+    If no filter is specified, all ops that implement `BufferizableOpInterface`
+    are bufferized. Ops from the `std` dialect are an exception: These ops are
+    always ignored, even if no filter is specified. When specifying a dialect
+    filter and `allow-unknown-ops` is not turned on, bufferization would fail
+    when encountering an op that is not included in the filter (even if it is
+    bufferizable).
+
+    For testing/debugging purposes, `test-analysis-only=1 print-conflicts=1`
+    prints analysis results and explains why an OpOperand was decided to
+    bufferize out-of-place. This is useful for understanding why One-Shot
+    Bufferize chose to insert a certain buffer copy.
+  }];
+  let options = [
+    Option<"allowReturnMemref", "allow-return-memref", "bool",
+            /*default=*/"false",
+           "Allows the return of memrefs (for testing purposes only)">,
+    Option<"allowUnknownOps", "allow-unknown-ops", "bool",
+           /*default=*/"false",
+           "Allows unknown (not bufferizable) ops in the input IR.">,
+    Option<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
+           /*default=*/"0",
+           "Test only: Analyze ops in random order with a given seed (fuzzer)">,
+    Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
+           "Specify if buffers should be deallocated. For compatibility with "
+           "core bufferization passes.">,
+    ListOption<"dialectFilter", "dialect-filter", "std::string",
+               "Restrict bufferization to ops from these dialects.",
+               "llvm::cl::MiscFlags::CommaSeparated">,
+    Option<"fullyDynamicLayoutMaps", "fully-dynamic-layout-maps", "bool",
+           /*default=*/"true",
+           "Generate MemRef types with dynamic offset+strides by default.">,
+    Option<"testAnalysisOnly", "test-analysis-only", "bool",
+            /*default=*/"false",
+           "Test only: Only run inplaceability analysis and annotate IR">,
+    Option<"printConflicts", "print-conflicts", "bool",
+            /*default=*/"false",
+           "Test only: Annotate IR with RaW conflicts. Requires "
+           "test-analysis-only.">,
+  ];
+  let constructor = "mlir::bufferization::createOneShotBufferizePass()";
+}
+
 def PromoteBuffersToStack : Pass<"promote-buffers-to-stack", "FuncOp"> {
   let summary = "Promotes heap-based allocations to automatically managed "
                 "stack-based allocations";
   let description = [{
     This pass implements a simple algorithm to convert heap-based memory
     allocations to stack-based ones. It uses a built-in heuristic to decide
     whether it makes sense to convert an allocation. Furthermore, dynamic
     shaped buffers that are limited by the rank of the tensor can be

mlir/include/mlir/InitAllDialects.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INITALLDIALECTS_H_
 #define MLIR_INITALLDIALECTS_H_
 
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arithmetic/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
+#include "mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Dialect.h"
 
 namespace mlir {
 
 /// Add all the MLIR dialects to the provided registry.
 inline void registerAllDialects(DialectRegistry &registry) {
   // clang-format off
registry.insert<acc::OpenACCDialect,
                   NVVM::NVVMDialect,
                   ROCDL::ROCDLDialect,
                   shape::ShapeDialect,
                   sparse_tensor::SparseTensorDialect,
                   tensor::TensorDialect,
                   tosa::TosaDialect,
                   x86vector::X86VectorDialect>();
   // clang-format on
+  arith::registerBufferizableOpInterfaceExternalModels(registry);
+  linalg::registerBufferizableOpInterfaceExternalModels(registry);
+  scf::registerBufferizableOpInterfaceExternalModels(registry);
+  tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingOpInterfaceExternalModels(registry);
+  vector::registerBufferizableOpInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.
 inline void registerAllDialects(MLIRContext &context) {
   DialectRegistry registry;
   registerAllDialects(registry);
   context.appendDialectRegistry(registry);
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

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

 //===- Bufferize.cpp - Bufferization utilities ----------------------------===//
 //
 // 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/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
 //===----------------------------------------------------------------------===//
 // BufferizeTypeConverter
 //===----------------------------------------------------------------------===//
 
void runOnOperation() override {
     // function.
     target.markUnknownOpDynamicallyLegal(
         [&](Operation *op) { return typeConverter.isLegal(op); });
 
     if (failed(applyFullConversion(func, target, std::move(patterns))))
       signalPassFailure();
   }
 };
+
+struct OneShotBufferizePass
+    : public OneShotBufferizeBase<OneShotBufferizePass> {
+  using OneShotBufferizeBase<OneShotBufferizePass>::OneShotBufferizeBase;
+
+  explicit OneShotBufferizePass(const AnalysisBufferizationOptions &options)
+      : options(options) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<bufferization::BufferizationDialect>();
+  }
+
+  void runOnOperation() override {
+    AnalysisBufferizationOptions opt;
+    if (!options) {
+      // Make new bufferization options if none were provided when creating the
+      // pass.
+      opt.allowReturnMemref = allowReturnMemref;
+      opt.allowUnknownOps = allowUnknownOps;
+      opt.analysisFuzzerSeed = analysisFuzzerSeed;
+      opt.createDeallocs = createDeallocs;
+      opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
+      opt.printConflicts = printConflicts;
+      opt.testAnalysisOnly = testAnalysisOnly;
+
+      BufferizationOptions::OpFilterEntry::FilterFn filterFn =
+          [&](Operation *op) {
+            // Disallow non-std dialect ops. I.e., no ops related to function
+            // calls.
+            if (op->getDialect()->getNamespace() ==
+                StandardOpsDialect::getDialectNamespace())
+              return false;
+            // Filter may be specified via options.
+            if (this->dialectFilter.hasValue())
+              return llvm::find(this->dialectFilter,
+                                op->getDialect()->getNamespace()) !=
+                     this->dialectFilter.end();
+            // No filter specified: All other ops are allowed.
+            return true;
+          };
+      opt.allowOperationInFilter(filterFn);
+    } else {
+      opt = *options;
+    }
+
+    ModuleOp moduleOp = getOperation();
+    if (failed(runOneShotBufferize(moduleOp, opt))) {
+      signalPassFailure();
+      return;
+    }
+
+    if (opt.testAnalysisOnly)
+      return;
+
+    OpPassManager cleanupPipeline("builtin.module");
+    cleanupPipeline.addPass(createCanonicalizerPass());
+    cleanupPipeline.addPass(createCSEPass());
+    cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
+    (void)runPipeline(cleanupPipeline, moduleOp);
+  }
+
+private:
+  llvm::Optional<AnalysisBufferizationOptions> options;
+};
 } // namespace
 
+std::unique_ptr<Pass> mlir::bufferization::createOneShotBufferizePass() {
+  return std::make_unique<OneShotBufferizePass>();
+}
+
+std::unique_ptr<Pass> mlir::bufferization::createOneShotBufferizePass(
+    const AnalysisBufferizationOptions &options) {
+  return std::make_unique<OneShotBufferizePass>(options);
+}
+
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::bufferization::createFinalizingBufferizePass() {
   return std::make_unique<FinalizingBufferizePass>();
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizableOpInterface-based Bufferization
 //===----------------------------------------------------------------------===//

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

LogicalResult bufferization::analyzeOp(Operation *op,
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, state);
 
   return success();
 }
 
 LogicalResult bufferization::runOneShotBufferize(
-    Operation *op, std::unique_ptr<AnalysisBufferizationOptions> options) {
-  AnalysisBufferizationState state(op, *options);
+    Operation *op, const AnalysisBufferizationOptions &options) {
+  AnalysisBufferizationState state(op, options);
   if (failed(analyzeOp(op, state)))
     return failure();
-  if (options->testAnalysisOnly)
+  if (options.testAnalysisOnly)
     return success();
   return bufferizeOp(op, state);
 }

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-compat.mlir

 // RUN: mlir-opt %s \
-// RUN:     -test-comprehensive-function-bufferize="allow-return-memref allow-unknown-ops create-deallocs=0" \
+// RUN:     -one-shot-bufferize="allow-return-memref allow-unknown-ops create-deallocs=0" \
 // RUN:     -split-input-file | \
 // RUN: FileCheck %s --check-prefix=CHECK-NODEALLOC
 
 // RUN: mlir-opt %s \
-// RUN:     -test-comprehensive-function-bufferize="allow-return-memref allow-unknown-ops create-deallocs=0" \
+// RUN:     -one-shot-bufferize="allow-return-memref allow-unknown-ops create-deallocs=0" \
 // RUN:     -buffer-deallocation | \
 // RUN: FileCheck %s --check-prefix=CHECK-BUFFERDEALLOC
 
 // CHECK-NODEALLOC-LABEL: func @out_of_place_bufferization
 // CHECK-BUFFERDEALLOC-LABEL: func @out_of_place_bufferization
 func @out_of_place_bufferization(%t1 : tensor<?xf32>) -> (f32, f32) {
   //     CHECK-NODEALLOC: memref.alloc
   //     CHECK-NODEALLOC: memref.copy

mlir/test/Dialect/Bufferization/Transforms/one-shot-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
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref allow-unknown-ops fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
 
 // Run fuzzer with different seeds.
-// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
-// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
-// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
-// RUN: mlir-opt %s -allow-unregistered-dialect -test-comprehensive-function-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
-// RUN: mlir-opt %s -allow-unregistered-dialect -test-comprehensive-function-bufferize="dialect-filter=scf allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-SCF
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=scf allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-SCF
 
 // CHECK: #[[$MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL: func @use_of_unknown_op_1(
-//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32, #[[$MAP]]>
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
 // CHECK-NO-LAYOUT-MAP-LABEL: func @use_of_unknown_op_1(
-//  CHECK-NO-LAYOUT-MAP-SAME:     %[[m1:.*]]: memref<?xf32>)
-func @use_of_unknown_op_1(%t1: tensor<?xf32> {linalg.inplaceable = true})
+//  CHECK-NO-LAYOUT-MAP-SAME:     %[[t1:.*]]: tensor<?xf32>
+func @use_of_unknown_op_1(%t1: tensor<?xf32>)
     -> vector<5xf32> {
   // ToTensorOp is generated because the function is bufferized and has a
   // memref block argument.
-  // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]] : memref<?xf32, #[[$MAP]]>
-  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
-  // CHECK-NO-LAYOUT-MAP: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]] : memref<?xf32>
-  // CHECK-NO-LAYOUT-MAP: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
+  // CHECK-NO-LAYOUT-MAP: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
   %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
 
   %idx = arith.constant 0 : index
   %cst = arith.constant 0.0 : f32
   // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]] : memref<?xf32, #[[$MAP]]>
   // CHECK: vector.transfer_read %[[dummy_memref]][%{{.*}}], %{{.*}} : memref<?xf32, #[[$MAP]]>
   // CHECK-NO-LAYOUT-MAP: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]] : memref<?xf32>
   // CHECK-NO-LAYOUT-MAP: vector.transfer_read %[[dummy_memref]][%{{.*}}], %{{.*}} : memref<?xf32>
   %1 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
   return %1 : vector<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @use_of_unknown_op_2(
-//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
-func @use_of_unknown_op_2(%t1: tensor<?xf32> {linalg.inplaceable = true})
-    -> tensor<?xf32> {
-  // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
-
-  // CHECK: %[[dummy1:.*]] = "test.dummy_op"(%[[m1_tensor]])
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
+func @use_of_unknown_op_2(%t1: tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK: %[[dummy1:.*]] = "test.dummy_op"(%[[t1]])
   %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
   // CHECK: %[[dummy2:.*]] = "test.another_dummy_op"(%[[dummy1]])
   %1 = "test.another_dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
 
-  // CHECK: %[[dummy2_memref:.*]] = bufferization.to_memref %[[dummy2]]
-  // CHECK: return %[[dummy2_memref]]
+  // CHECK: return %[[dummy2]]
   return %1 : tensor<?xf32>
 }
 
 // -----
 
+// CHECK: #[[$MAP2:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
 // CHECK-LABEL: func @use_of_unknown_op_3(
-//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
-func @use_of_unknown_op_3(%t1: tensor<?xf32> {linalg.inplaceable = true})
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
+func @use_of_unknown_op_3(%t1: tensor<?xf32>)
     -> (vector<5xf32>, vector<5xf32>) {
   %idx = arith.constant 0 : index
   %cst = arith.constant 0.0 : f32
-  // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]]
   // CHECK: %[[v1:.*]] = vector.transfer_read %[[m1]]
   %1 = vector.transfer_read %t1[%idx], %cst : tensor<?xf32>, vector<5xf32>
 
-  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
   %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
-  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
+  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]] : memref<?xf32, #[[$MAP2]]>
   // CHECK: %[[v2:.*]] = vector.transfer_read %[[dummy_memref]]
   %2 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
 
   // CHECK: return %[[v1]], %[[v2]]
   return %1, %2 : vector<5xf32>, vector<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @use_of_unknown_op_4(
-//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
-func @use_of_unknown_op_4(%t1: tensor<?xf32> {linalg.inplaceable = true})
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
+func @use_of_unknown_op_4(%t1: tensor<?xf32>)
     -> (vector<5xf32>, tensor<?xf32>) {
   %idx = arith.constant 0 : index
   %cst = arith.constant 0.0 : f32
 
-  // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
-  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
   %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
 
   // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
   // CHECK: %[[v1:.*]] = vector.transfer_read %[[dummy_memref]]
   %1 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
 
   // CHECK: %[[another_dummy:.*]] = "test.another_dummy_op"(%[[dummy]])
   %2 = "test.another_dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
 
-  // CHECK: %[[another_dummy_memref:.*]] = bufferization.to_memref %[[another_dummy]]
-  // CHECK: return %[[v1]], %[[another_dummy_memref]]
+  // CHECK: return %[[v1]], %[[another_dummy]]
   return %1, %2 : vector<5xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @use_of_bufferizable_op_in_unbufferizable_op
-//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
 func @use_of_bufferizable_op_in_unbufferizable_op(
     %t1: tensor<?xf32>, %o: index, %s: index) -> (tensor<?xf32>, tensor<?xf32>) {
+  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]]
   // CHECK: %[[subview:.*]] = memref.subview %[[m1]]
   %0 = tensor.extract_slice %t1[%o][%s][1] : tensor<?xf32> to tensor<?xf32>
   // CHECK: %[[subview_tensor:.*]] = bufferization.to_tensor %[[subview]]
   // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[subview_tensor]])
   %1 = "test.dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
-  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
-  // CHECK: return %[[subview]], %[[dummy_memref]]
+  // CHECK: return %[[subview_tensor]], %[[dummy]]
   return %0, %1 : tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @unused_unknown_op(
-//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
 func @unused_unknown_op(%t1 : tensor<?xf32>) -> vector<5xf32> {
   %idx = arith.constant 0 : index
   %cst = arith.constant 0.0 : f32
-  // ToTensorOp is inserted to pass in the result of the above bufferized op.
-  // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+
+  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]]
   // CHECK: vector.transfer_read %[[m1]]
   %1 = vector.transfer_read %t1[%idx], %cst : tensor<?xf32>, vector<5xf32>
 
-  // CHECK: "test.dummy_op"(%[[m1_tensor]])
+  // CHECK: "test.dummy_op"(%[[t1]])
   "test.dummy_op"(%t1) : (tensor<?xf32>) -> ()
 
   return %1 : vector<5xf32>
 }
 
 // -----
 
+// CHECK: #[[$MAP3:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+// 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: %[[m1:.*]] = memref.alloc() {{.*}} : memref<10xf32>
+  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10xf32>
+  // CHECK: %[[alloc_casted:.*]] = memref.cast %[[alloc]] : memref<10xf32> to memref<10xf32, #[[$MAP3]]>
+  // CHECK: %[[m1_casted:.*]] = memref.cast %[[m1]] : memref<10xf32> to memref<10xf32, #[[$MAP3]]>
+  %t1 = linalg.init_tensor [10] : tensor<10xf32>
+
+  // CHECK: linalg.fill(%{{.*}}, %[[m1]])
+  // CHECK: %[[filled_tensor:.*]] = bufferization.to_tensor %[[m1_casted]]
+  %filled = linalg.fill(%cst, %t1) : f32, tensor<10xf32> -> tensor<10xf32>
+
+  // The transfer_write is out-of-place because "dummy_op" may read.
+  // CHECK: memref.copy %[[m1]], %[[alloc]]
+  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
+  // CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc_casted]]
+  %1 = vector.transfer_write %v, %filled[%idx] : vector<5xf32>, tensor<10xf32>
+
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[filled_tensor]])
+  %2 = "test.dummy_op"(%filled) : (tensor<10xf32>) -> (tensor<10xf32>)
+
+  // CHECK: memref.dealloc %[[alloc]]
+  // CHECK: memref.dealloc %[[m1]]
+  // CHECK: return %[[alloc_tensor]], %[[dummy]]
+  return %1, %2 : tensor<10xf32>, tensor<10xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @unknown_op_not_writable
-//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
 func @unknown_op_not_writable(
     %t1 : tensor<?xf32>, %v :  vector<5xf32>, %idx : index) -> tensor<?xf32> {
-  // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
-  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
   // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
   %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> (tensor<?xf32>)
 
   // The result of an unknown op is not writable. Always generate a copy.
-  // Note: This copy is essential for partial bufferization. Otherwise, we could
-  // introducing a RaW conflict.
   // CHECK: %[[dim:.*]] = tensor.dim %[[dummy]]
   // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
   // CHECK: memref.copy %[[dummy_memref]], %[[alloc]]
   // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
   %1 = vector.transfer_write %v, %0[%idx] : vector<5xf32>, tensor<?xf32>
 
-  // CHECK: return %[[alloc]]
+  // CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc]]
+  // CHECK: return %[[alloc_tensor]]
   return %1 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-TENSOR-LABEL: func @simple_tensor_test(
 //  CHECK-TENSOR-SAME:     %[[t1:.*]]: tensor<?xf32>
 func @simple_tensor_test(%t1 : tensor<?xf32>, %f : f32) -> tensor<?xf32> {

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

-// RUN: mlir-opt %s -test-comprehensive-function-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
-// RUN: mlir-opt %s -test-comprehensive-function-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
-// RUN: mlir-opt %s -test-comprehensive-function-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
-// RUN: mlir-opt %s -test-comprehensive-function-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 // CHECK-LABEL: func @use_tensor_func_arg(
 //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
 func @use_tensor_func_arg(%A : tensor<?xf32>) -> (vector<4xf32>) {
   %c0 = arith.constant 0 : index
   %f0 = arith.constant 0.0 : f32
 
   // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]

 // CHECK-LABEL: func private @private_func
 func private @private_func(tensor<?xf32>) -> ()
 
 // CHECK-LABEL: func @empty_func()
 func @empty_func() -> () {
   return
 }
 
-// -----
 
-// CHECK-LABEL: func @rank_reducing
-func @rank_reducing(
-    %i: index, %j: index,
-    %arg0: tensor<8x18x32xf32>) 
-      -> tensor<?x1x6x8xf32> {
-  %c1 = arith.constant 1 : index
-  %c6 = arith.constant 6 : index
-  %c8 = arith.constant 8 : index
-  %c32 = arith.constant 32 : index
-  %c0 = arith.constant 0 : index
-  %0 = linalg.init_tensor [4, 1, 6, 8] : tensor<4x1x6x8xf32>
-  %1 = tensor.cast %0 : tensor<4x1x6x8xf32> to tensor<?x1x6x8xf32>
-  %2 = linalg.init_tensor [1, 6, 8] : tensor<1x6x8xf32>
-  %5 = scf.for %arg7 = %c0 to %c32 step %c8 iter_args(%arg8 = %1) -> (tensor<?x1x6x8xf32>) {
-    %7 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg7)
-    %8 = tensor.extract_slice %arg0[%i, %j, %arg7] [1, 6, 8] [1, 1, 1] : tensor<8x18x32xf32> to tensor<1x6x8xf32>
-    %9 = scf.for %arg9 = %c0 to %c6 step %c1 iter_args(%arg10 = %2) -> (tensor<1x6x8xf32>) {
-      %11 = tensor.extract_slice %8[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x6x8xf32> to tensor<1x1x8xf32>
-      %12 = tensor.insert_slice %11 into %arg10[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x1x8xf32> into tensor<1x6x8xf32>
-      scf.yield %12 : tensor<1x6x8xf32>
-    }
-    %10 = tensor.insert_slice %9 into %arg8[%7, 0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] : tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
-    scf.yield %10 : tensor<?x1x6x8xf32>
-  }
-  return %5: tensor<?x1x6x8xf32>
-}

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

func @write_after_select_read_one(
   %w = tensor.insert %cst into %s[%idx] : tensor<?xf32>
 
   // CHECK: %[[f:.*]] = memref.load %[[t1]]
   %f = tensor.extract %t1[%idx] : tensor<?xf32>
 
   // CHECK: return %[[f]], %[[select]]
   return %f, %w : f32, tensor<?xf32>
 }
+
+// -----
+
+// A regression test to make sure that we handle rank-reducing extract_slice
+// correctly.
+
+// CHECK-LABEL: func @rank_reducing
+func @rank_reducing(
+    %i: index, %j: index,
+    %arg0: tensor<8x18x32xf32>)
+      -> tensor<?x1x6x8xf32> {
+  %c1 = arith.constant 1 : index
+  %c6 = arith.constant 6 : index
+  %c8 = arith.constant 8 : index
+  %c32 = arith.constant 32 : index
+  %c0 = arith.constant 0 : index
+  %0 = linalg.init_tensor [4, 1, 6, 8] : tensor<4x1x6x8xf32>
+  %1 = tensor.cast %0 : tensor<4x1x6x8xf32> to tensor<?x1x6x8xf32>
+  %2 = linalg.init_tensor [1, 6, 8] : tensor<1x6x8xf32>
+  %5 = scf.for %arg7 = %c0 to %c32 step %c8 iter_args(%arg8 = %1) -> (tensor<?x1x6x8xf32>) {
+    %7 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg7)
+    %8 = tensor.extract_slice %arg0[%i, %j, %arg7] [1, 6, 8] [1, 1, 1] : tensor<8x18x32xf32> to tensor<1x6x8xf32>
+    %9 = scf.for %arg9 = %c0 to %c6 step %c1 iter_args(%arg10 = %2) -> (tensor<1x6x8xf32>) {
+      %11 = tensor.extract_slice %8[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x6x8xf32> to tensor<1x1x8xf32>
+      %12 = tensor.insert_slice %11 into %arg10[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x1x8xf32> into tensor<1x6x8xf32>
+      scf.yield %12 : tensor<1x6x8xf32>
+    }
+    %10 = tensor.insert_slice %9 into %arg8[%7, 0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] : tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
+    scf.yield %10 : tensor<?x1x6x8xf32>
+  }
+  return %5: tensor<?x1x6x8xf32>
+}

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

-//===- TestComprehensiveBufferize.cpp - Test Comprehensive Bufferize ------===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-//
-// This file implements logic for testing Comprehensive Bufferize.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.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/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/Linalg/Passes.h"
-#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
-#include "mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h"
-#include "mlir/Dialect/StandardOps/IR/Ops.h"
-#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
-#include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
-#include "mlir/Pass/PassManager.h"
-#include "mlir/Transforms/Passes.h"
-
-using namespace mlir;
-using namespace mlir::linalg;
-using namespace mlir::linalg::comprehensive_bufferize;
-using namespace mlir::bufferization;
-
-namespace {
-/// A helper struct for FunctionBufferize and ModuleBufferize. Both passes are
-/// mostly identical.
-struct TestComprehensiveFunctionBufferize
-    : public PassWrapper<TestComprehensiveFunctionBufferize,
-                         OperationPass<FuncOp>> {
-  StringRef getArgument() const final {
-    return "test-comprehensive-function-bufferize";
-  }
-
-  StringRef getDescription() const final {
-    return "Test Comprehensive Bufferize of FuncOps (body only).";
-  }
-
-  TestComprehensiveFunctionBufferize() = default;
-  TestComprehensiveFunctionBufferize(
-      const TestComprehensiveFunctionBufferize &pass)
-      : PassWrapper(pass) {}
-
-  void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
-                    memref::MemRefDialect, tensor::TensorDialect,
-                    vector::VectorDialect, scf::SCFDialect, StandardOpsDialect,
-                    arith::ArithmeticDialect, AffineDialect>();
-    affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
-    arith::registerBufferizableOpInterfaceExternalModels(registry);
-    linalg::registerBufferizableOpInterfaceExternalModels(registry);
-    scf::registerBufferizableOpInterfaceExternalModels(registry);
-    tensor::registerBufferizableOpInterfaceExternalModels(registry);
-    vector::registerBufferizableOpInterfaceExternalModels(registry);
-  }
-
-  void runOnOperation() override;
-
-  Option<bool> allowReturnMemref{
-      *this, "allow-return-memref",
-      llvm::cl::desc("Allow returning/yielding memrefs from functions/blocks"),
-      llvm::cl::init(false)};
-  Option<bool> allowUnknownOps{
-      *this, "allow-unknown-ops",
-      llvm::cl::desc(
-          "Allows the return of memrefs (for testing purposes only)"),
-      llvm::cl::init(false)};
-  Option<bool> testAnalysisOnly{
-      *this, "test-analysis-only",
-      llvm::cl::desc(
-          "Only runs inplaceability analysis (for testing purposes only)"),
-      llvm::cl::init(false)};
-  Option<unsigned> analysisFuzzerSeed{
-      *this, "analysis-fuzzer-seed",
-      llvm::cl::desc("Analyze ops in random order with a given seed (fuzzer)"),
-      llvm::cl::init(0)};
-  ListOption<std::string> dialectFilter{
-      *this, "dialect-filter",
-      llvm::cl::desc("Bufferize only ops from the specified dialects"),
-      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
-  Option<bool> fullyDynamicLayoutMaps{
-      *this, "fully-dynamic-layout-maps",
-      llvm::cl::desc("Use fully dynamic layout maps on memref types"),
-      llvm::cl::init(true)};
-  Option<bool> createDeallocs{
-      *this, "create-deallocs",
-      llvm::cl::desc("Specify if buffers should be deallocated"),
-      llvm::cl::init(true)};
-};
-} // namespace
-
-void TestComprehensiveFunctionBufferize::runOnOperation() {
-  auto options = std::make_unique<AnalysisBufferizationOptions>();
-  options->allowReturnMemref = allowReturnMemref;
-  options->allowUnknownOps = allowUnknownOps;
-  options->testAnalysisOnly = testAnalysisOnly;
-  options->analysisFuzzerSeed = analysisFuzzerSeed;
-  options->fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
-  options->createDeallocs = createDeallocs;
-
-  if (dialectFilter.hasValue()) {
-    options->hasFilter = true;
-    for (const std::string &dialectNamespace : dialectFilter)
-      options->allowDialectInFilter(dialectNamespace);
-  }
-
-  Operation *op = getOperation();
-  if (failed(runOneShotBufferize(op, std::move(options))))
-    return;
-
-  if (testAnalysisOnly)
-    return;
-
-  OpPassManager cleanupPipeline("builtin.func");
-  cleanupPipeline.addPass(createCanonicalizerPass());
-  cleanupPipeline.addPass(createCSEPass());
-  cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
-  (void)this->runPipeline(cleanupPipeline, op);
-}
-
-namespace mlir {
-namespace test {
-void registerTestComprehensiveFunctionBufferize() {
-  PassRegistration<TestComprehensiveFunctionBufferize>();
-}
-} // namespace test
-} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerInliner();
 void registerMemRefBoundCheck();
 void registerPatternsTestPass();
 void registerSimpleParametricTilingPass();
 void registerTestAffineLoopParametricTilingPass();
 void registerTestAliasAnalysisPass();
 void registerTestBuiltinAttributeInterfaces();
 void registerTestCallGraphPass();
-void registerTestComprehensiveFunctionBufferize();
 void registerTestConstantFold();
 void registerTestGpuSerializeToCubinPass();
 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
void registerTestPasses() {
   mlir::test::registerTestConstantFold();
   mlir::test::registerTestDiagnosticsPass();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToHsacoPass();
 #endif
-  mlir::test::registerTestComprehensiveFunctionBufferize();
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandTanhPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestGpuParallelLoopMappingPass();
   mlir::test::registerTestIRVisitorsPass();