[Mlir-commits] [mlir] c1fef4e - [mlir][bufferization] Make `TensorCopyInsertionPass` a test pass
Matthias Springer
llvmlistbot at llvm.org
Fri Dec 2 06:39:33 PST 2022
Author: Matthias Springer
Date: 2022-12-02T15:38:02+01:00
New Revision: c1fef4e88ad461724f7454d7d8177f7fe30cf58c
URL: https://github.com/llvm/llvm-project/commit/c1fef4e88ad461724f7454d7d8177f7fe30cf58c
DIFF: https://github.com/llvm/llvm-project/commit/c1fef4e88ad461724f7454d7d8177f7fe30cf58c.diff
LOG: [mlir][bufferization] Make `TensorCopyInsertionPass` a test pass
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 `SparsificationAndBufferizationPass` 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.
Differential Revision: https://reviews.llvm.org/D138915
Added:
mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
mlir/test/lib/Dialect/Bufferization/CMakeLists.txt
mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp
Modified:
mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp
mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt
mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-invalid.mlir
mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space.mlir
mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir
mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir
mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir
mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir
mlir/test/Dialect/Tensor/one-shot-bufferize-tensor-copy-insertion.mlir
mlir/test/lib/Dialect/CMakeLists.txt
mlir/tools/mlir-opt/CMakeLists.txt
mlir/tools/mlir-opt/mlir-opt.cpp
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
Removed:
mlir/lib/Dialect/SparseTensor/Transforms/DenseBufferizationPass.cpp
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
index 52ba73b60181a..972fa2b0c49d0 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
@@ -97,11 +97,6 @@ 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
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index ebad5bb4772f1..cbf01e4f788ea 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -340,37 +340,6 @@ def PromoteBuffersToStack : Pass<"promote-buffers-to-stack", "func::FuncOp"> {
];
}
-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
index 9b04c376f5c2c..252463566c414 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -162,8 +162,11 @@ createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true,
// Other rewriting rules and passes.
//===----------------------------------------------------------------------===//
-std::unique_ptr<Pass> createDenseBufferizationPass(
- const bufferization::OneShotBufferizationOptions &options);
+std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
+ 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
index 69ca30ade1b6a..677f0d5a7f4be 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp
@@ -161,45 +161,3 @@ mlir::bufferization::insertTensorCopies(Operation *op,
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
index 8c9cd7d085470..6697285594366 100644
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -52,25 +52,12 @@ getBufferizationOptions(bool analysisOnly) {
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(createSparsificationAndBufferizationPass(
+ 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
index 00c624fd6e085..a44d2117f5891 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt
@@ -1,14 +1,14 @@
add_mlir_dialect_library(MLIRSparseTensorTransforms
BufferizableOpInterfaceImpl.cpp
CodegenUtils.cpp
- DenseBufferizationPass.cpp
- Sparsification.cpp
SparseBufferRewriting.cpp
SparseTensorCodegen.cpp
SparseTensorConversion.cpp
SparseTensorPasses.cpp
SparseTensorRewriting.cpp
SparseVectorization.cpp
+ Sparsification.cpp
+ SparsificationAndBufferizationPass.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/DenseBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/DenseBufferizationPass.cpp
deleted file mode 100644
index 49a882f31dccd..0000000000000
--- 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/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
new file mode 100644
index 0000000000000..b3d10813407c6
--- /dev/null
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -0,0 +1,155 @@
+//===- SparsificationAndBufferizationPass.cpp - Tensor to Memref Lowering -===//
+//
+// 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 lowers tensor ops to memref ops, regardless of whether they are
+/// dense or sparse.
+///
+/// One-Shot Analysis is used to detect RaW conflicts and to insert buffer
+/// copies of the tensor level (`insertTensorCopies`). Afterwards, the lowering
+/// of tensor ops to memref ops follows a
diff erent code path depending on
+/// whether the op is sparse or dense:
+///
+/// * Sparse tensor ops are lowered through Sparsification and follow-up pass
+/// that lowers sparse_tensor dialect ops.
+/// * Dense tensor ops are lowered through BufferizableOpInterface
+/// implementations.
+class SparsificationAndBufferizationPass
+ : public PassWrapper<SparsificationAndBufferizationPass,
+ OperationPass<ModuleOp>> {
+public:
+ SparsificationAndBufferizationPass(
+ const bufferization::OneShotBufferizationOptions &bufferizationOptions,
+ const SparsificationOptions &sparsificationOptions,
+ const SparseTensorConversionOptions &sparseTensorConversionOptions,
+ bool enableRuntimeLibrary, bool enableBufferInitialization)
+ : bufferizationOptions(bufferizationOptions),
+ sparsificationOptions(sparsificationOptions),
+ sparseTensorConversionOptions(sparseTensorConversionOptions),
+ enableRuntimeLibrary(enableRuntimeLibrary),
+ enableBufferInitialization(enableBufferInitialization) {}
+
+ /// Bufferize all dense ops. This assumes that no further analysis is needed
+ /// and that all required buffer copies were already inserted by
+ /// `insertTensorCopies` in the form of `bufferization.alloc_tensor` ops.
+ LogicalResult runDenseBufferization() {
+ 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;
+ });
+ return bufferization::bufferizeOp(getOperation(), bufferizationOptions,
+ /*copyBeforeWrite=*/false,
+ &denseOpFilter);
+ }
+
+ void runOnOperation() override {
+ {
+ // Run enabling transformations.
+ OpPassManager pm("builtin.module");
+ pm.addPass(createPreSparsificationRewritePass());
+ if (failed(runPipeline(pm, getOperation())))
+ return signalPassFailure();
+ }
+
+ // Insert tensor copies. This step runs One-Shot Analysis (which analyzes
+ // SSA use-def chains of tensor IR) and decides where buffer copies are
+ // needed and where buffers can be written to in-place. These decisions are
+ // materialized in the IR in the form of `bufferization.alloc_tensor` ops.
+ //
+ // Note: All following steps in this pass must be careful not to modify the
+ // structure of the IR (i.e., tensor use-def chains), as that could
+ // invalidate the results of the analysis. From now on, only small and
+ // localized rewrites are allowed, such as replacing a tensor op with its
+ // memref equivalent.
+ if (failed(bufferization::insertTensorCopies(getOperation(),
+ bufferizationOptions)))
+ return signalPassFailure();
+
+ // `testAnalysisOnly` is a debug/testing flag. If set, the results of
+ // OneShotAnalysis are added to the IR via attributes. In that case, do not
+ // continue with the remaining pipeline.
+ if (bufferizationOptions.testAnalysisOnly)
+ return;
+
+ // Bufferize all sparse ops. No further analysis is needed. All required
+ // buffer copies were already inserted by `insertTensorCopies` in the form
+ // of `bufferization.alloc_tensor` 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.
+ if (failed(runDenseBufferization()))
+ signalPassFailure();
+ }
+
+private:
+ bufferization::OneShotBufferizationOptions bufferizationOptions;
+ SparsificationOptions sparsificationOptions;
+ SparseTensorConversionOptions sparseTensorConversionOptions;
+ bool enableRuntimeLibrary;
+ bool enableBufferInitialization;
+};
+} // namespace sparse_tensor
+} // namespace mlir
+
+std::unique_ptr<Pass> mlir::createSparsificationAndBufferizationPass(
+ const bufferization::OneShotBufferizationOptions &bufferizationOptions,
+ const SparsificationOptions &sparsificationOptions,
+ const SparseTensorConversionOptions &sparseTensorConversionOptions,
+ bool enableRuntimeLibrary, bool enableBufferInitialization) {
+ return std::make_unique<
+ mlir::sparse_tensor::SparsificationAndBufferizationPass>(
+ 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
index ba71090ec1ba8..b9ea67b1b7e5a 100644
--- 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
index 1397ccd1b0484..7679350d47fa4 100644
--- 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
index 2d3a1f1bc006b..4dea1cbcfaa4f 100644
--- 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
index 17dca3faf2e9c..d231a5cc46e3e 100644
--- 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
index e518efb98f378..68a275f775c39 100644
--- 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
index de3104540af85..e0c3f08df9ce1 100755
--- 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
index d4732976025fb..84fe150400d70 100644
--- 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
index 0000000000000..a183d02cefed7
--- /dev/null
+++ b/mlir/test/lib/Dialect/Bufferization/CMakeLists.txt
@@ -0,0 +1,12 @@
+# Exclude tests from libMLIR.so
+add_mlir_library(MLIRBufferizationTestPasses
+ 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
index 0000000000000..6c4011d462a77
--- /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
index 0c98d0c6959d8..11c223620b58c 100644
--- 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/CMakeLists.txt b/mlir/tools/mlir-opt/CMakeLists.txt
index 011656cf4fadd..ca33fae2d4e48 100644
--- a/mlir/tools/mlir-opt/CMakeLists.txt
+++ b/mlir/tools/mlir-opt/CMakeLists.txt
@@ -15,6 +15,7 @@ if(MLIR_INCLUDE_TESTS)
MLIRTestFuncToLLVM
MLIRAffineTransformsTestPasses
MLIRArithTestPasses
+ MLIRBufferizationTestPasses
MLIRDLTITestPasses
MLIRFuncTestPasses
MLIRGPUTestPasses
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index cfce28ffb9e5a..3295ad22dbaa7 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -113,6 +113,7 @@ void registerTestRecursiveTypesPass();
void registerTestSCFUtilsPass();
void registerTestShapeMappingPass();
void registerTestSliceAnalysisPass();
+void registerTestTensorCopyInsertionPass();
void registerTestTensorTransforms();
void registerTestTilingInterface();
void registerTestTopologicalSortAnalysisPass();
@@ -216,6 +217,7 @@ void registerTestPasses() {
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
index fbd2cf6bca783..88749e23032b5 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -6865,6 +6865,7 @@ cc_binary(
"//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
index 0ea2bbc267800..091ebb4855c3c 100644
--- 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(
],
)
+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 = [
More information about the Mlir-commits
mailing list