[Mlir-commits] [mlir] 8154494 - [mlir][sparse] refactor sparsification and bufferization pass into proper TD pass
Aart Bik
llvmlistbot at llvm.org
Thu Aug 17 15:45:11 PDT 2023
Author: Aart Bik
Date: 2023-08-17T15:45:03-07:00
New Revision: 8154494e28364f3eb85f2bba624f33224e4aed08
URL: https://github.com/llvm/llvm-project/commit/8154494e28364f3eb85f2bba624f33224e4aed08
DIFF: https://github.com/llvm/llvm-project/commit/8154494e28364f3eb85f2bba624f33224e4aed08.diff
LOG: [mlir][sparse] refactor sparsification and bufferization pass into proper TD pass
Registering the SparsificationAndBufferization into a proper TD pass
has the advantage that it can be invoked and tested in isolation. This
change also moves some bufferization specific set up from the pipeline
file into the pass file, keeping the logic more locally.
Reviewed By: Peiming
Differential Revision: https://reviews.llvm.org/D158219
Added:
Modified:
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 2b237842bc6ed5..c1e217675020f0 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -22,6 +22,7 @@
//===----------------------------------------------------------------------===//
namespace mlir {
+
namespace bufferization {
struct OneShotBufferizationOptions;
} // namespace bufferization
@@ -215,12 +216,13 @@ std::unique_ptr<Pass> createStorageSpecifierToLLVMPass();
//===----------------------------------------------------------------------===//
// The mini-pipeline for sparsification and bufferization.
-//
-// Note that this mini-pipeline is not defined through the tablegen pass
-// mechanism, and, thus, is not individually available through the command-line.
-// It is solely used as part of the full sparse compiler pipeline.
//===----------------------------------------------------------------------===//
+bufferization::OneShotBufferizationOptions
+getBufferizationOptionsForSparsification(bool analysisOnly);
+
+std::unique_ptr<Pass> createSparsificationAndBufferizationPass();
+
std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index f59785c0fcef9a..ab7fffac88d928 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -373,4 +373,23 @@ def StorageSpecifierToLLVM : Pass<"sparse-storage-specifier-to-llvm", "ModuleOp"
];
}
+def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "ModuleOp"> {
+ let summary = "Mini-pipeline that combines bufferization and sparsifiation";
+ let description = [{
+ This pass forms a mini-pipeline that combines bufferization and sparsifiation.
+ }];
+ let constructor = "mlir::createSparsificationAndBufferizationPass()";
+ let dependentDialects = [
+ "affine::AffineDialect",
+ "arith::ArithDialect",
+ "bufferization::BufferizationDialect",
+ "gpu::GPUDialect",
+ "LLVM::LLVMDialect",
+ "linalg::LinalgDialect",
+ "memref::MemRefDialect",
+ "scf::SCFDialect",
+ "sparse_tensor::SparseTensorDialect",
+ ];
+}
+
#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES
diff --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
index 990f8f7327d809..b962f7ae9b4eea 100644
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -25,31 +25,6 @@
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/Passes.h"
-using namespace mlir;
-using namespace mlir::sparse_tensor;
-
-/// Return configuration options for One-Shot Bufferize.
-static bufferization::OneShotBufferizationOptions
-getBufferizationOptions(bool analysisOnly) {
- using namespace bufferization;
- OneShotBufferizationOptions options;
- options.bufferizeFunctionBoundaries = true;
- // TODO(springerm): To spot memory leaks more easily, returning dense allocs
- // should be disallowed.
- options.allowReturnAllocs = true;
- options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
- options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
- const BufferizationOptions &options) {
- return getMemRefTypeWithStaticIdentityLayout(
- cast<TensorType>(value.getType()), memorySpace);
- };
- if (analysisOnly) {
- options.testAnalysisOnly = true;
- options.printConflicts = true;
- }
- return options;
-}
-
//===----------------------------------------------------------------------===//
// Pipeline implementation.
//===----------------------------------------------------------------------===//
@@ -58,7 +33,8 @@ void mlir::sparse_tensor::buildSparseCompiler(
OpPassManager &pm, const SparseCompilerOptions &options) {
pm.addNestedPass<func::FuncOp>(createLinalgGeneralizationPass());
pm.addPass(createSparsificationAndBufferizationPass(
- getBufferizationOptions(options.testBufferizationAnalysisOnly),
+ getBufferizationOptionsForSparsification(
+ options.testBufferizationAnalysisOnly),
options.sparsificationOptions(), options.sparseTensorConversionOptions(),
options.createSparseDeallocs, options.enableRuntimeLibrary,
options.enableBufferInitialization, options.vectorLength,
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index fbbf6c18316fd3..9af525d55bf34e 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -8,6 +8,7 @@
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
@@ -18,15 +19,21 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/Passes.h"
using namespace mlir;
-using namespace mlir::func;
namespace mlir {
+
+#define GEN_PASS_DEF_SPARSIFICATIONANDBUFFERIZATION
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
+
namespace sparse_tensor {
/// Return `true` if one of the given types is a sparse tensor type.
@@ -50,8 +57,8 @@ static bool containsSparseTensor(TypeRange types) {
/// * Dense tensor ops are lowered through BufferizableOpInterface
/// implementations.
class SparsificationAndBufferizationPass
- : public PassWrapper<SparsificationAndBufferizationPass,
- OperationPass<ModuleOp>> {
+ : public impl::SparsificationAndBufferizationBase<
+ SparsificationAndBufferizationPass> {
public:
SparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
@@ -97,12 +104,6 @@ class SparsificationAndBufferizationPass
return success();
}
- void getDependentDialects(::mlir::DialectRegistry ®istry) const override {
- registry.insert<bufferization::BufferizationDialect>();
- registry.insert<gpu::GPUDialect>();
- registry.insert<LLVM::LLVMDialect>();
- }
-
void runOnOperation() override {
{
// Run enabling transformations.
@@ -179,7 +180,42 @@ class SparsificationAndBufferizationPass
} // namespace sparse_tensor
} // namespace mlir
-std::unique_ptr<Pass> mlir::createSparsificationAndBufferizationPass(
+mlir::bufferization::OneShotBufferizationOptions
+mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
+ using namespace mlir::bufferization;
+ OneShotBufferizationOptions options;
+ options.bufferizeFunctionBoundaries = true;
+ // TODO(springerm): To spot memory leaks more easily, returning dense allocs
+ // should be disallowed.
+ options.allowReturnAllocs = true;
+ options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
+ options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
+ const BufferizationOptions &options) {
+ return getMemRefTypeWithStaticIdentityLayout(
+ cast<TensorType>(value.getType()), memorySpace);
+ };
+ if (analysisOnly) {
+ options.testAnalysisOnly = true;
+ options.printConflicts = true;
+ }
+ return options;
+}
+
+std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass() {
+ SparsificationOptions sparseOptions;
+ SparseTensorConversionOptions convOptions;
+ return createSparsificationAndBufferizationPass(
+ getBufferizationOptionsForSparsification(/*analysisOnly=*/false),
+ sparseOptions, convOptions,
+ /*createSparseDeallocs=*/false,
+ /*enableRuntimeLibrary=*/false,
+ /*enableBufferInitialization=*/false,
+ /*vectorLength=*/0,
+ /*enableVLAVectorization=*/false,
+ /*enableSIMDIndex32=*/false);
+}
+
+std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions,
More information about the Mlir-commits
mailing list