diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
@@ -97,11 +97,6 @@
/// 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
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -340,37 +340,6 @@
];
}
-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 = [{
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -162,8 +162,11 @@
// 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);
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp b/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp
--- a/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp
@@ -161,45 +161,3 @@
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 ®istry) 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);
-}
diff --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -52,25 +52,13 @@
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());
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt
@@ -1,8 +1,8 @@
add_mlir_dialect_library(MLIRSparseTensorTransforms
BufferizableOpInterfaceImpl.cpp
CodegenUtils.cpp
- DenseBufferizationPass.cpp
Sparsification.cpp
+ SparseBufferizationPass.cpp
SparseBufferRewriting.cpp
SparseTensorCodegen.cpp
SparseTensorConversion.cpp
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/DenseBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/DenseBufferizationPass.cpp
deleted file mode 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/DenseBufferizationPass.cpp
+++ /dev/null
@@ -1,73 +0,0 @@
-//===- 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);
-}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp
new file mode 100644
--- /dev/null
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferizationPass.cpp
@@ -0,0 +1,110 @@
+//===- 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);
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-invalid.mlir
--- a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-invalid.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-invalid.mlir
@@ -1,4 +1,4 @@
-// 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.
diff --git a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space.mlir b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space.mlir
--- a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space.mlir
@@ -1,4 +1,4 @@
-// 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>) {
diff --git a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir
--- a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir
@@ -1,6 +1,6 @@
-// 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>
diff --git a/mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir
--- a/mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir
+++ b/mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir
@@ -1,5 +1,5 @@
-// 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>
diff --git a/mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir b/mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir
--- a/mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir
+++ b/mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir
@@ -1,5 +1,5 @@
-// 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" ],
diff --git a/mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir b/mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir
--- a/mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir
@@ -1,4 +1,4 @@
-// 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" ] }>
diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize-tensor-copy-insertion.mlir b/mlir/test/Dialect/Tensor/one-shot-bufferize-tensor-copy-insertion.mlir
--- a/mlir/test/Dialect/Tensor/one-shot-bufferize-tensor-copy-insertion.mlir
+++ b/mlir/test/Dialect/Tensor/one-shot-bufferize-tensor-copy-insertion.mlir
@@ -1,5 +1,5 @@
-// 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>
diff --git a/mlir/test/lib/Dialect/Bufferization/CMakeLists.txt b/mlir/test/lib/Dialect/Bufferization/CMakeLists.txt
new file mode 100644
--- /dev/null
+++ b/mlir/test/lib/Dialect/Bufferization/CMakeLists.txt
@@ -0,0 +1,12 @@
+# Exclude tests from libMLIR.so
+add_mlir_library(MLIRArithTestPasses
+ TestTensorCopyInsertion.cpp
+
+ EXCLUDE_FROM_LIBMLIR
+
+ LINK_LIBS PUBLIC
+ MLIRBufferizationDialect
+ MLIRBufferizationTransforms
+ MLIRIR
+ MLIRPass
+)
diff --git a/mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp b/mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp
new file mode 100644
--- /dev/null
+++ b/mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp
@@ -0,0 +1,78 @@
+//===- 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 ®istry) 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
diff --git a/mlir/test/lib/Dialect/CMakeLists.txt b/mlir/test/lib/Dialect/CMakeLists.txt
--- a/mlir/test/lib/Dialect/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/CMakeLists.txt
@@ -1,5 +1,6 @@
add_subdirectory(Affine)
add_subdirectory(Arith)
+add_subdirectory(Bufferization)
add_subdirectory(DLTI)
add_subdirectory(Func)
add_subdirectory(GPU)
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -113,6 +113,7 @@
void registerTestSCFUtilsPass();
void registerTestShapeMappingPass();
void registerTestSliceAnalysisPass();
+void registerTestTensorCopyInsertionPass();
void registerTestTensorTransforms();
void registerTestTilingInterface();
void registerTestTopologicalSortAnalysisPass();
@@ -216,6 +217,7 @@
mlir::test::registerTestSCFUtilsPass();
mlir::test::registerTestShapeMappingPass();
mlir::test::registerTestSliceAnalysisPass();
+ mlir::test::registerTestTensorCopyInsertionPass();
mlir::test::registerTestTensorTransforms();
mlir::test::registerTestTilingInterface();
mlir::test::registerTestTopologicalSortAnalysisPass();
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -6866,6 +6866,7 @@
"//mlir/test:TestAffine",
"//mlir/test:TestAnalysis",
"//mlir/test:TestArith",
+ "//mlir/test:TestBufferization",
"//mlir/test:TestDLTI",
"//mlir/test:TestDialect",
"//mlir/test:TestFunc",
diff --git a/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
--- a/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
@@ -701,6 +701,19 @@
],
)
+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 = [