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 @@ -22,6 +22,7 @@ //===----------------------------------------------------------------------===// namespace mlir { + namespace bufferization { struct OneShotBufferizationOptions; } // namespace bufferization @@ -215,12 +216,13 @@ //===----------------------------------------------------------------------===// // 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 createSparsificationAndBufferizationPass(); + std::unique_ptr 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 --- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td @@ -373,4 +373,23 @@ ]; } +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 --- 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(value.getType()), memorySpace); - }; - if (analysisOnly) { - options.testAnalysisOnly = true; - options.printConflicts = true; - } - return options; -} - //===----------------------------------------------------------------------===// // Pipeline implementation. //===----------------------------------------------------------------------===// @@ -58,7 +33,8 @@ OpPassManager &pm, const SparseCompilerOptions &options) { pm.addNestedPass(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 --- 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 @@ /// * Dense tensor ops are lowered through BufferizableOpInterface /// implementations. class SparsificationAndBufferizationPass - : public PassWrapper> { + : public impl::SparsificationAndBufferizationBase< + SparsificationAndBufferizationPass> { public: SparsificationAndBufferizationPass( const bufferization::OneShotBufferizationOptions &bufferizationOptions, @@ -97,12 +104,6 @@ return success(); } - void getDependentDialects(::mlir::DialectRegistry ®istry) const override { - registry.insert(); - registry.insert(); - registry.insert(); - } - void runOnOperation() override { { // Run enabling transformations. @@ -179,7 +180,42 @@ } // namespace sparse_tensor } // namespace mlir -std::unique_ptr 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(value.getType()), memorySpace); + }; + if (analysisOnly) { + options.testAnalysisOnly = true; + options.printConflicts = true; + } + return options; +} + +std::unique_ptr 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::createSparsificationAndBufferizationPass( const bufferization::OneShotBufferizationOptions &bufferizationOptions, const SparsificationOptions &sparsificationOptions, const SparseTensorConversionOptions &sparseTensorConversionOptions,