[mlir][bufferization][WIP] Make `TensorCopyInsertionPass` a test pass

[springerm]

Summary

TensorCopyInsertion should not have been exposed as a pass. This was a flaw in the original design. It is a preparation step for bufferization and certain transforms (that would otherwise be legal) are illegal between TensorCopyInsertion and actual rewrite to MemRef ops. Therefore, even if broken down as two separate steps internally, they should be exposed as a single pass.

This change affects the sparse compiler, which uses TensorCopyInsertionPass. A new SparseBufferizationPass is added to replace all passes in the sparse tensor pipeline from TensorCopyInsertionPass until the actual bufferization (rewrite to memref/non-tensor). It is generally unsafe to run arbitrary passes in-between, in particular passes that hoist tensor ops out of loops or change SSA use-def chains along tensor ops.

[jreiffers]

Thanks for the fix!

[springerm]

(Will wait for Aart's comments.)

[aartbik]

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

60

At very first glance, I am a bit -1 on moving these passes out of the pipeline into a bufferization pass.
I understand why you are doing this but is there any way you can still do this while not breaking this flow here?

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp

1

name seems copy paste

[Peiming]

Thanks Matthias, I believe that this is the correct thing to do!

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp

51

Is there any reason that we need to do insertTensorCopy before Sparsification in particular? If we can move all dense related bufferization after sparsification, wouldn't that be more clear? Or am I missing anything? It seems that it could be dangerous to do optimization based on simple SSA def-use chain after InsertTensorCopies, I am a little bit worry that it might cause other issues in the future (or there is other remaining issue in the current pipeline).

If we somehow must perform insertTensorCopy here, what types of optimization are sound afterwards (i.e., what assumption can we hold for the dense tensor SSA value)? is there any general rules that we should follow?

60

nit: If we want to wrap these passes into the same class, I'd prefer including createPreSparsificationRewritePass in this pass as well, so that Pre- and Post-rewriting passes are in the same place.

[springerm]

The requirements are:

• TensorCopyInsertion must run before sparsification, sparsetensorconversion, densebufferization (otherwise we may be missing copies).
• The same pass that runs TensorCopyInsertion should also rewrite all tensor type ops to memref type ops. I.e. all tensor type ops are gone afterwards.

I wanted to keep the passes in the original pipeline but couldn’t find a way. Do you see a better layering?

[springerm]

Is there any reason that we need to do insertTensorCopy before Sparsification

sparsification removes the linalg.generic with loops right? memref loops maybe even? we cannot analyze these anymore. tensorcopyinsertion must run before introducing any memrefs.

It seems that it could be dangerous to do optimization based on simple SSA def-use chain after InsertTensorCopies

no more analyses happen after tensorcopyinsertion. in particular no bufferization analysis

If we somehow must perform insertTensorCopy here, what types of optimization are sound afterwards (i.e., what assumption can we hold for the dense tensor SSA value)? is there any general rules that we should follow?

the short answer is none. ideally do not touch the IR apart from rewriting tensor ops to memref ops. the dense bufferization does a Operation::walk over the IR instead of using the greedypatternrewriter for that reason. the greedypatternrewriter could call folding patterns which are unsafe. basically only BufferizableOpInterface::bufferize (and it’s sparse equivalent) is allowed

[aartbik]

Thanks for finding a solution, Matthias. Some feedback below, but assuming no two-pass solution can be made, I would at least like to see this documented and using a better "pass" name for the compound pass.

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp

1

That was perhaps a bit vague, but I mean that after the renaming that should be

"SparseBufferizationPass.cpp"

34

I gave this another good look, and grouping this all together is probably a good way to make sure stuff happens in the right order and with the right information.

I guess I objected a little on creating a "mini"pipeline as part of the larger pipeline (or perhaps a "super"pass on top of our other passes). If you still see a way of putting the two parts you need in two passes and keep the original pipeline, that would have my preference. But guessing that is not feasible, I think my other objection came from the name SparseBufferizationPass.

I think if (1) we give this a better name, like SparsificationPass (which is heavily overloaded by now, so I am very open to better names) (2) we document really well that this is grouping stuff together to form a logical consistent unit (that does all the rewriting and bufferization "atomically" and cannot be split in clean passes), I am okay with this change.

60

+1 on keeping the pre/post together

[springerm]

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp

60

PreSparsificationRewriting contains a pattern that should not run after TensorCopyInsertion: FoldInvariantYield. This was the reason for one of the miscompiles that we were seeing.

Can we rename PreSparsificationRewriting and PostSparsificationRewriting into something more descriptive (that ideally describes what the pass is actually doing)?

[springerm]

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp

34

Imo SparsificationPass does not describe it very accurately because the pass is doing more than that. This pass basically converts a tensor program into a memref program. In particular, it also handles dense tensor IR. But this is just naming and for you to decide.

Some ideas that came to mind:

• sparse_tensor::SparseAndDenseBufferizationPass
• sparse_tensor::BufferizationPass
• sparse_tensor::ConvertTensorToMemrefPass
• sparse_tensor::SparsificationPass and we mention in the pass description that it also does other stuff.

If we go with SparsificationPass: There is already a SparsificationPass, so we would have to rename that one. I always thought of it as ConvertSparseLinalgToLoopsPass. Because there is also a ConvertLinalgToLoopsPass which does the same thing, just for the dense case.

92–96

Do you see a better way than having all these options here as a field? I wanted to use SparseCompilerOptions, but it is defined in Pipelines/Passes.h (not in Transforms). I guess I could create a new SparseBufferizationOptions (or whatever we call the pass).

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp

60

It should still be okay to put PreSparsificationRewriting into the class (before TensorCopyInsertion)? But I think you are right about renaming, probably that is a better way to make thing more clear.

I think FoldInvariantYield is to enable FuseSparseMulitplyOverAdd, maybe we can just call it FuseSparseOperations?
I do not have a good name for PostSparsificationRewriting though.

@aartbik
What do you think?

[aartbik]

Okay, this is LGTM provided that
(1) we name the new "mini"pipeline, sparse_tensor::SparsificationAndBufferizationPass
(2) we move PreSparsificationRewriting into that pass, running first, but just so they are all logically contained
(3) we keep the Pre/PostSparsificationRewriting names; nothing is perfect but these names reflect the current function reasonably well

[wrengr]

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

60–61

I wonder if this is still relevant here? Before this differential this conditional bailed out before doing all our various sparse passes, but now it's only doing so after all of SparseBufferizationPass. So I'm thinking the conditional should either be moved into SparseBufferizationPass or else removed entirely

[springerm]

address comments

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

 
 /// Create a pass that tries to eliminate tensor.empty ops that are anchored on
 /// insert_slice ops.
 std::unique_ptr<Pass> createEmptyTensorEliminationPass();
 
 /// Create a pass that bufferizes ops from the bufferization dialect.
 std::unique_ptr<Pass> createBufferizationBufferizePass();
 
-/// Create a pass that resolves out-of-place tensor OpOperands with copies.
-std::unique_ptr<Pass> createTensorCopyInsertionPass();
-std::unique_ptr<Pass>
-createTensorCopyInsertionPass(const OneShotBufferizationOptions &options);
-
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Register external models for AllocationOpInterface.
 void registerAllocationOpInterfaceExternalModels(DialectRegistry &registry);
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H

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

Option<"maxAllocSizeInBytes", "max-alloc-size-in-bytes", "unsigned",
            /*default=*/"1024",
            "Maximal size in bytes to promote allocations to stack.">,
     Option<"maxRankOfAllocatedMemRef", "max-rank-of-allocated-memref", "unsigned",
            /*default=*/"1",
            "Maximal memref rank to promote dynamic buffers.">,
   ];
 }
 
-def TensorCopyInsertion : Pass<"tensor-copy-insertion"> {
-  let summary = "Make all tensor IR inplaceable by inserting copies";
-  let description = [{
-    This pass runs One-Shot Analysis and inserts copies for all OpOperands that
-    were decided to bufferize out-of-place. After running this pass, a
-    bufferization can write to buffers directly (without making copies) and no
-    longer has to care about potential read-after-write conflicts.
-
-    Note: By default, all newly inserted tensor copies/allocs (i.e., newly
-    created `bufferization.alloc_tensor` ops) that do not escape block are
-    annotated with `escape = false`. If `create-allocs` is unset, all newly
-    inserted tensor copies/allocs are annotated with `escape = true`. In that
-    case, they are not getting deallocated when bufferizing the IR.
-  }];
-  let options = [
-    Option<"allowReturnAllocs", "allow-return-allocs", "bool",
-            /*default=*/"false",
-           "Allows returning/yielding new allocations from a block.">,
-    Option<"bufferizeFunctionBoundaries", "bufferize-function-boundaries",
-           "bool", /*default=*/"0",
-           "Bufferize function boundaries (experimental).">,
-    Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
-           "Specify if new allocations should be deallocated.">,
-    Option<"mustInferMemorySpace", "must-infer-memory-space", "bool",
-           /*default=*/"false",
-           "The memory space of an memref types must always be inferred. If "
-           "unset, a default memory space of 0 is used otherwise.">,
-  ];
-  let constructor = "mlir::bufferization::createTensorCopyInsertionPass()";
-}
-
 def EmptyTensorElimination : Pass<"eliminate-empty-tensors"> {
   let summary = "Try to eliminate all tensor.empty ops.";
   let description = [{
     This pass tries to eliminate all insert_slice op-anchored tensor.empty ops.
     I.e., when a value that is equivalent to an tensor.empty op is inserted into
     another tensor, this pass tries to rewrite the IR in such a way that the
     destination tensor of the insert_slice op is used directly instead of the
     tensor.empty result.
   }];
   let constructor = "mlir::bufferization::createEmptyTensorEliminationPass()";
 }
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES

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

 std::unique_ptr<Pass>
 createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true,
                                     bool enableConvert = true);
 
 //===----------------------------------------------------------------------===//
 // Other rewriting rules and passes.
 //===----------------------------------------------------------------------===//
 
-std::unique_ptr<Pass> createDenseBufferizationPass(
-    const bufferization::OneShotBufferizationOptions &options);
+std::unique_ptr<Pass> createSparseBufferizationPass(
+    const bufferization::OneShotBufferizationOptions &bufferizationOptions,
+    const SparsificationOptions &sparsificationOptions,
+    const SparseTensorConversionOptions &sparseTensorConversionOptions,
+    bool enableRuntimeLibrary, bool enableBufferInitialization);
 
 void populateSparseBufferRewriting(RewritePatternSet &patterns,
                                    bool enableBufferInitialization);
 
 std::unique_ptr<Pass> createSparseBufferRewritePass();
 std::unique_ptr<Pass>
 createSparseBufferRewritePass(bool enableBufferInitialization);
 

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

WalkResult result = op->walk([&](Operation *op) {
     if (failed(bufferizableOp.resolveConflicts(rewriter, state)))
       return WalkResult::interrupt();
 
     return WalkResult::advance();
   });
 
   return failure(result.wasInterrupted());
 }
-
-namespace {
-struct TensorCopyInsertionPass
-    : public bufferization::impl::TensorCopyInsertionBase<
-          TensorCopyInsertionPass> {
-  TensorCopyInsertionPass() : options(llvm::None) {}
-  TensorCopyInsertionPass(const OneShotBufferizationOptions &options)
-      : options(options) {}
-
-  void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<bufferization::BufferizationDialect>();
-  }
-
-  void runOnOperation() override {
-    if (options) {
-      if (failed(insertTensorCopies(getOperation(), *options)))
-        signalPassFailure();
-    } else {
-      OneShotBufferizationOptions options;
-      options.allowReturnAllocs = allowReturnAllocs;
-      options.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
-      options.createDeallocs = createDeallocs;
-      if (mustInferMemorySpace)
-        options.defaultMemorySpace = None;
-      if (failed(insertTensorCopies(getOperation(), options)))
-        signalPassFailure();
-    }
-  }
-
-private:
-  Optional<OneShotBufferizationOptions> options;
-};
-} // namespace
-
-std::unique_ptr<Pass> mlir::bufferization::createTensorCopyInsertionPass() {
-  return std::make_unique<TensorCopyInsertionPass>();
-}
-
-std::unique_ptr<Pass> mlir::bufferization::createTensorCopyInsertionPass(
-    const OneShotBufferizationOptions &options) {
-  return std::make_unique<TensorCopyInsertionPass>(options);
-}

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

 
 //===----------------------------------------------------------------------===//
 // Pipeline implementation.
 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   pm.addNestedPass<func::FuncOp>(createLinalgGeneralizationPass());
-  pm.addPass(
-      bufferization::createTensorCopyInsertionPass(getBufferizationOptions(
-          /*analysisOnly=*/options.testBufferizationAnalysisOnly)));
+  pm.addPass(createPreSparsificationRewritePass());
+  pm.addPass(createSparseBufferizationPass(
+      getBufferizationOptions(options.testBufferizationAnalysisOnly),
+      options.sparsificationOptions(), options.sparseTensorConversionOptions(),
+      options.enableRuntimeLibrary, options.enableBufferInitialization));
   if (options.testBufferizationAnalysisOnly)
     return;
-  pm.addPass(createPreSparsificationRewritePass());
-  pm.addPass(createSparsificationPass(options.sparsificationOptions()));
-  pm.addPass(createPostSparsificationRewritePass(options.enableRuntimeLibrary));
-  if (options.enableRuntimeLibrary) {
-    pm.addPass(createSparseTensorConversionPass(
-        options.sparseTensorConversionOptions()));
-  } else {
-    pm.addPass(
-        createSparseTensorCodegenPass(options.enableBufferInitialization));
-    pm.addPass(
-        createSparseBufferRewritePass(options.enableBufferInitialization));
-  }
-  pm.addPass(createDenseBufferizationPass(
-      getBufferizationOptions(/*analysisOnly=*/false)));
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());
   // TODO(springerm): Add sparse support to the BufferDeallocation pass and add
   // it to this pipeline.
   pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());
   pm.addNestedPass<func::FuncOp>(createConvertVectorToSCFPass());
   pm.addNestedPass<func::FuncOp>(createConvertSCFToCFPass());

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

 add_mlir_dialect_library(MLIRSparseTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   CodegenUtils.cpp
-  DenseBufferizationPass.cpp
   Sparsification.cpp
+  SparseBufferizationPass.cpp
   SparseBufferRewriting.cpp
   SparseTensorCodegen.cpp
   SparseTensorConversion.cpp
   SparseTensorPasses.cpp
   SparseTensorRewriting.cpp
   SparseVectorization.cpp
 
   ADDITIONAL_HEADER_DIRS

mlir/lib/Dialect/SparseTensor/Transforms/DenseBufferizationPass.cpp

-//===- DenseBufferizationPass.cpp - Dense bufferization pass --------------===//
-//
-// 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/SparseTensor/Transforms/Passes.h"
-
-#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
-#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
-#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
-
-using namespace mlir;
-using namespace mlir::func;
-
-namespace mlir {
-namespace sparse_tensor {
-
-/// Return `true` if one of the given types is a sparse tensor type.
-static bool containsSparseTensor(TypeRange types) {
-  for (Type t : types)
-    if (getSparseTensorEncoding(t))
-      return true;
-  return false;
-}
-
-/// A pass that bufferizes only dense tensor ops and ignores all sparse tensor
-/// ops. No buffer copies are inserted. All tensor OpOperands must be
-/// inplacable.
-class BufferizeDenseOpsPass
-    : public PassWrapper<BufferizeDenseOpsPass, OperationPass<ModuleOp>> {
-public:
-  BufferizeDenseOpsPass(
-      const bufferization::OneShotBufferizationOptions &options)
-      : options(options) {}
-
-  void runOnOperation() override {
-    // Disallow all sparse tensor ops, so that only dense tensor ops are
-    // bufferized.
-    bufferization::OpFilter opFilter;
-    opFilter.allowOperation([&](Operation *op) {
-      if (containsSparseTensor(TypeRange(op->getResults())) ||
-          containsSparseTensor(TypeRange(op->getOperands())))
-        return false;
-      if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
-        FunctionType funcType = funcOp.getFunctionType();
-        if (containsSparseTensor(funcType.getInputs()) ||
-            containsSparseTensor(funcType.getResults()))
-          return false;
-      }
-      return true;
-    });
-
-    if (failed(bufferization::bufferizeOp(getOperation(), options,
-                                          /*copyBeforeWrite=*/false,
-                                          &opFilter)))
-      signalPassFailure();
-  }
-
-private:
-  bufferization::OneShotBufferizationOptions options;
-};
-} // namespace sparse_tensor
-} // namespace mlir
-
-std::unique_ptr<Pass> mlir::createDenseBufferizationPass(
-    const bufferization::OneShotBufferizationOptions &options) {
-  return std::make_unique<mlir::sparse_tensor::BufferizeDenseOpsPass>(options);
-}

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp

+//===- DenseBufferizationPass.cpp - Dense bufferization pass --------------===//
+//
+// 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/SparseTensor/Transforms/Passes.h"
+
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
+#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
+#include "mlir/Pass/PassManager.h"
+
+using namespace mlir;
+using namespace mlir::func;
+
+namespace mlir {
+namespace sparse_tensor {
+
+/// Return `true` if one of the given types is a sparse tensor type.
+static bool containsSparseTensor(TypeRange types) {
+  for (Type t : types)
+    if (getSparseTensorEncoding(t))
+      return true;
+  return false;
+}
+
+/// A pass that bufferizes dense and sparse ops.
+class SparseBufferizationPass
+    : public PassWrapper<SparseBufferizationPass, OperationPass<ModuleOp>> {
+public:
+  SparseBufferizationPass(
+      const bufferization::OneShotBufferizationOptions &bufferizationOptions,
+      const SparsificationOptions &sparsificationOptions,
+      const SparseTensorConversionOptions &sparseTensorConversionOptions,
+      bool enableRuntimeLibrary, bool enableBufferInitialization)
+      : bufferizationOptions(bufferizationOptions),
+        sparsificationOptions(sparsificationOptions),
+        sparseTensorConversionOptions(sparseTensorConversionOptions),
+        enableRuntimeLibrary(enableRuntimeLibrary),
+        enableBufferInitialization(enableBufferInitialization) {}
+
+  void runOnOperation() override {
+    // Insert tensor copies.
+    if (failed(bufferization::insertTensorCopies(getOperation(),
+                                                 bufferizationOptions)))
+      return signalPassFailure();
+    if (bufferizationOptions.testAnalysisOnly)
+      return;
+
+    // Bufferize all sparse ops.
+    OpPassManager pm("builtin.module");
+    pm.addPass(createSparsificationPass(sparsificationOptions));
+    pm.addPass(createPostSparsificationRewritePass(enableRuntimeLibrary));
+    if (enableRuntimeLibrary) {
+      pm.addPass(
+          createSparseTensorConversionPass(sparseTensorConversionOptions));
+    } else {
+      pm.addPass(createSparseTensorCodegenPass(enableBufferInitialization));
+      pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
+    }
+    if (failed(runPipeline(pm, getOperation())))
+      return signalPassFailure();
+
+    // Bufferize all dense ops.
+    bufferization::OpFilter denseOpFilter;
+    denseOpFilter.allowOperation([&](Operation *op) {
+      if (containsSparseTensor(TypeRange(op->getResults())) ||
+          containsSparseTensor(TypeRange(op->getOperands())))
+        return false;
+      if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
+        FunctionType funcType = funcOp.getFunctionType();
+        if (containsSparseTensor(funcType.getInputs()) ||
+            containsSparseTensor(funcType.getResults()))
+          return false;
+      }
+      return true;
+    });
+    if (failed(bufferization::bufferizeOp(getOperation(), bufferizationOptions,
+                                          /*copyBeforeWrite=*/false,
+                                          &denseOpFilter)))
+      return signalPassFailure();
+  }
+
+private:
+  bufferization::OneShotBufferizationOptions bufferizationOptions;
+  SparsificationOptions sparsificationOptions;
+  SparseTensorConversionOptions sparseTensorConversionOptions;
+  bool enableRuntimeLibrary;
+  bool enableBufferInitialization;
+};
+} // namespace sparse_tensor
+} // namespace mlir
+
+std::unique_ptr<Pass> mlir::createSparseBufferizationPass(
+    const bufferization::OneShotBufferizationOptions &bufferizationOptions,
+    const SparsificationOptions &sparsificationOptions,
+    const SparseTensorConversionOptions &sparseTensorConversionOptions,
+    bool enableRuntimeLibrary, bool enableBufferInitialization) {
+  return std::make_unique<mlir::sparse_tensor::SparseBufferizationPass>(
+      bufferizationOptions, sparsificationOptions,
+      sparseTensorConversionOptions, enableRuntimeLibrary,
+      enableBufferInitialization);
+}

mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-invalid.mlir

-// RUN: mlir-opt %s -tensor-copy-insertion="must-infer-memory-space" -split-input-file -verify-diagnostics
+// RUN: mlir-opt %s -test-tensor-copy-insertion="must-infer-memory-space" -split-input-file -verify-diagnostics
 
 // An alloc is inserted but the copy is emitted. Therefore, the memory space
 // should be specified on the alloc_tensor op.
 func.func @memory_space_of_unknown_op() -> (tensor<10xf32>, tensor<10xf32>) {
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.0 : f32
   // expected-error @+1 {{could not infer memory space}}
   %t = bufferization.alloc_tensor() : tensor<10xf32>
   %s = tensor.insert %cst into %t[%c0] : tensor<10xf32>
   return %s, %t : tensor<10xf32>, tensor<10xf32>
 }

mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space.mlir

-// RUN: mlir-opt %s -tensor-copy-insertion="must-infer-memory-space" -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-tensor-copy-insertion="must-infer-memory-space" -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @unknown_op_copy
 func.func @unknown_op_copy() -> (tensor<10xf32>, tensor<10xf32>) {
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.0 : f32
   // CHECK: %[[dummy:.*]] = "test.dummy_op"() : () -> tensor<10xf32>
   %t = "test.dummy_op"() : () -> tensor<10xf32>
   // CHECK: %[[copy:.*]] = bufferization.alloc_tensor() copy(%[[dummy]]) {bufferization.escape = [false]} : tensor<10xf32>

mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir

-// RUN: mlir-opt %s -tensor-copy-insertion -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
-// RUN: mlir-opt %s -tensor-copy-insertion="create-deallocs=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-DEALLOC
+// RUN: mlir-opt %s -test-tensor-copy-insertion -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
+// RUN: mlir-opt %s -test-tensor-copy-insertion="create-deallocs=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-DEALLOC
 
 // CHECK-LABEL: func @read_after_write_conflict(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
 // CHECK-FUNC-LABEL: func @read_after_write_conflict(
 // CHECK-NO-DEALLOC-LABEL: func @read_after_write_conflict(
 func.func @read_after_write_conflict(%t: tensor<?xf32>, %idx: index, %f: f32)
   -> (tensor<?xf32>, tensor<?xf32>)
 {

mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir

-// RUN: mlir-opt %s -tensor-copy-insertion="allow-return-allocs" -allow-unregistered-dialect -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
+// RUN: mlir-opt %s -test-tensor-copy-insertion="allow-return-allocs" -allow-unregistered-dialect -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
 
 // CHECK-LABEL: func @scf_for(
 //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>, %[[B:.*]]: tensor<?xf32>
 func.func @scf_for(%A : tensor<?xf32>, %B : tensor<?xf32>,
                    %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
   // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {bufferization.escape = [false]} : tensor<?xf32>

mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir

-// RUN: mlir-opt %s -tensor-copy-insertion="allow-return-allocs" | FileCheck %s
-// RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" | FileCheck %s --check-prefix=CHECK-FUNC
+// RUN: mlir-opt %s -test-tensor-copy-insertion="allow-return-allocs" | FileCheck %s
+// RUN: mlir-opt %s -test-tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" | FileCheck %s --check-prefix=CHECK-FUNC
 
 #DCSR = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ],
   dimOrdering = affine_map<(i,j) -> (i,j)>
 }>
 
 // CHECK-LABEL: func @bufferization_alloc_tensor
 // CHECK-FUNC-LABEL: func @bufferization_alloc_tensor

mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir

-// RUN: mlir-opt %s  --tensor-copy-insertion --pre-sparsification-rewrite --sparsification --cse | FileCheck %s
+// RUN: mlir-opt %s  --test-tensor-copy-insertion --pre-sparsification-rewrite --sparsification --cse | FileCheck %s
 
 #SM = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 #trait_matmul = {
   indexing_maps = [
     affine_map<(d0, d1, d2) -> (d1, d0)>,
     affine_map<(d0, d1, d2) -> (d0, d2)>,
     affine_map<(d0, d1, d2) -> (d1, d2)>

mlir/test/Dialect/Tensor/one-shot-bufferize-tensor-copy-insertion.mlir

-// RUN: mlir-opt %s -tensor-copy-insertion -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
+// RUN: mlir-opt %s -test-tensor-copy-insertion -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
 
 // CHECK-LABEL: func @extract_slice(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
 // CHECK-FUNC-LABEL: func @extract_slice(
 func.func @extract_slice(%t: tensor<?xf32>, %idx: index, %f: f32)
   -> (tensor<5xf32>, tensor<?xf32>)
 {
   // CHECK: %[[extract_slice:.*]] = tensor.extract_slice %[[t]][10] [5] [1]
   %0 = tensor.extract_slice %t[10][5][1] : tensor<?xf32> to tensor<5xf32>
   // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[extract_slice]]) {bufferization.escape = [false]} : tensor<5xf32>
   // CHECK-FUNC: bufferization.alloc_tensor() copy(%{{.*}}) {bufferization.escape = [true]} : tensor<5xf32>
   // CHECK: %[[insert:.*]] = tensor.insert %{{.*}} into %[[alloc]]
   %1 = tensor.insert %f into %0[%idx] : tensor<5xf32>
   // CHECK: return %[[insert]], %[[t]]
   return %1, %t : tensor<5xf32>, tensor<?xf32>
 }

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

+# Exclude tests from libMLIR.so
+add_mlir_library(MLIRArithTestPasses
+  TestTensorCopyInsertion.cpp
+
+  EXCLUDE_FROM_LIBMLIR
+
+  LINK_LIBS PUBLIC
+  MLIRBufferizationDialect
+  MLIRBufferizationTransforms
+  MLIRIR
+  MLIRPass
+)

mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp

+//===- TestTensorCopyInsertion.cpp - Bufferization Analysis -----*- c++ -*-===//
+//
+// 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/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
+#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// This pass runs One-Shot Analysis and inserts copies for all OpOperands that
+/// were decided to bufferize out-of-place. After running this pass, a
+/// bufferization can write to buffers directly (without making copies) and no
+/// longer has to care about potential read-after-write conflicts.
+///
+/// Note: By default, all newly inserted tensor copies/allocs (i.e., newly
+/// created `bufferization.alloc_tensor` ops) that do not escape block are
+/// annotated with `escape = false`. If `create-allocs` is unset, all newly
+/// inserted tensor copies/allocs are annotated with `escape = true`. In that
+/// case, they are not getting deallocated when bufferizing the IR.
+struct TestTensorCopyInsertionPass
+    : public PassWrapper<TestTensorCopyInsertionPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTensorCopyInsertionPass)
+
+  TestTensorCopyInsertionPass() = default;
+  TestTensorCopyInsertionPass(const TestTensorCopyInsertionPass &pass)
+      : PassWrapper(pass) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<bufferization::BufferizationDialect>();
+  }
+  StringRef getArgument() const final { return "test-tensor-copy-insertion"; }
+  StringRef getDescription() const final {
+    return "Module pass to test Tensor Copy Insertion";
+  }
+
+  void runOnOperation() override {
+    bufferization::OneShotBufferizationOptions options;
+    options.allowReturnAllocs = allowReturnAllocs;
+    options.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
+    options.createDeallocs = createDeallocs;
+    if (mustInferMemorySpace)
+      options.defaultMemorySpace = None;
+    if (failed(bufferization::insertTensorCopies(getOperation(), options)))
+      signalPassFailure();
+  }
+
+  Option<bool> allowReturnAllocs{
+      *this, "allow-return-allocs",
+      llvm::cl::desc("Allows returning/yielding new allocations from a block."),
+      llvm::cl::init(false)};
+  Option<bool> bufferizeFunctionBoundaries{
+      *this, "bufferize-function-boundaries",
+      llvm::cl::desc("Bufferize function boundaries."), llvm::cl::init(false)};
+  Option<bool> createDeallocs{
+      *this, "create-deallocs",
+      llvm::cl::desc("Specify if new allocations should be deallocated."),
+      llvm::cl::init(true)};
+  Option<bool> mustInferMemorySpace{
+      *this, "must-infer-memory-space",
+      llvm::cl::desc(
+          "The memory space of an memref types must always be inferred. If "
+          "unset, a default memory space of 0 is used otherwise."),
+      llvm::cl::init(false)};
+};
+} // namespace
+
+namespace mlir::test {
+void registerTestTensorCopyInsertionPass() {
+  PassRegistration<TestTensorCopyInsertionPass>();
+}
+} // namespace mlir::test

mlir/test/lib/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
 add_subdirectory(Arith)
+add_subdirectory(Bufferization)
 add_subdirectory(DLTI)
 add_subdirectory(Func)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVM)
 add_subdirectory(Math)
 add_subdirectory(MemRef)
 add_subdirectory(NVGPU)

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

 void registerTestPadFusion();
 void registerTestPDLByteCodePass();
 void registerTestPDLLPasses();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
 void registerTestShapeMappingPass();
 void registerTestSliceAnalysisPass();
+void registerTestTensorCopyInsertionPass();
 void registerTestTensorTransforms();
 void registerTestTilingInterface();
 void registerTestTopologicalSortAnalysisPass();
 void registerTestTransformDialectEraseSchedulePass();
 void registerTestTransformDialectInterpreterPass();
 void registerTestVectorLowerings();
 void registerTestNvgpuLowerings();
 } // namespace test
#endif
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPadFusion();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestPDLLPasses();
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();
   mlir::test::registerTestShapeMappingPass();
   mlir::test::registerTestSliceAnalysisPass();
+  mlir::test::registerTestTensorCopyInsertionPass();
   mlir::test::registerTestTensorTransforms();
   mlir::test::registerTestTilingInterface();
   mlir::test::registerTestTopologicalSortAnalysisPass();
   mlir::test::registerTestTransformDialectEraseSchedulePass();
   mlir::test::registerTestTransformDialectInterpreterPass();
   mlir::test::registerTestVectorLowerings();
   mlir::test::registerTestNvgpuLowerings();
 }

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

deps = [
         ":SCFToGPU",
         ":Support",
         ":Transforms",
         "//llvm:AllTargetsCodeGens",
         "//llvm:Support",
         "//mlir/test:TestAffine",
         "//mlir/test:TestAnalysis",
         "//mlir/test:TestArith",
+        "//mlir/test:TestBufferization",
         "//mlir/test:TestDLTI",
         "//mlir/test:TestDialect",
         "//mlir/test:TestFunc",
         "//mlir/test:TestFuncToLLVM",
         "//mlir/test:TestGPU",
         "//mlir/test:TestIR",
         "//mlir/test:TestLinalg",
         "//mlir/test:TestLLVM",

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

deps = [
         "//mlir:MathDialect",
         "//mlir:Pass",
         "//mlir:Transforms",
         "//mlir:VectorDialect",
     ],
 )
 
 cc_library(
+    name = "TestBufferization",
+    srcs = glob(["lib/Dialect/Bufferization/*.cpp"]),
+    defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],
+    includes = ["lib/Dialect/Test"],
+    deps = [
+        "//mlir:BufferizationDialect",
+        "//mlir:BufferizationTransforms",
+        "//mlir:IR",
+        "//mlir:Pass",
+    ],
+)
+
+cc_library(
     name = "TestShapeDialect",
     srcs = [
         "lib/Dialect/Shape/TestShapeFunctions.cpp",
         "lib/Dialect/Shape/TestShapeMappingAnalysis.cpp",
     ],
     deps = [
         "//llvm:Support",
         "//mlir:FuncDialect",