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 @@ -8,10 +8,6 @@ // // This header file defines prototypes of all sparse tensor passes. // -// In general, this file takes the approach of keeping "mechanism" (the -// actual steps of applying a transformation) completely separate from -// "policy" (heuristics for when and where to apply transformations). -// //===----------------------------------------------------------------------===// #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_ @@ -21,15 +17,15 @@ #include "mlir/Pass/Pass.h" #include "mlir/Transforms/DialectConversion.h" +//===----------------------------------------------------------------------===// +// Include the generated pass header (which needs some early definitions). +//===----------------------------------------------------------------------===// + namespace mlir { namespace bufferization { struct OneShotBufferizationOptions; } // namespace bufferization -//===----------------------------------------------------------------------===// -// The Sparsification pass. -//===----------------------------------------------------------------------===// - /// Defines a parallelization strategy. Any independent loop is a candidate /// for parallelization. The loop is made parallel if (1) allowed by the /// strategy (e.g., AnyStorageOuterLoop considers either a dense or sparse @@ -43,12 +39,25 @@ kAnyStorageAnyLoop }; -// TODO : Zero copy is disabled due to correctness bugs.Tracker #64316 +/// Defines data movement strategy between host and device for GPU. +// TODO : Zero copy is disabled due to correctness bugs (tracker #64316) enum class GPUDataTransferStrategy { kRegularDMA, kZeroCopy, kPinnedDMA }; #define GEN_PASS_DECL #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc" +//===----------------------------------------------------------------------===// +// The PreSparsificationRewriting pass. +//===----------------------------------------------------------------------===// + +void populatePreSparsificationRewriting(RewritePatternSet &patterns); + +std::unique_ptr createPreSparsificationRewritePass(); + +//===----------------------------------------------------------------------===// +// The Sparsification pass. +//===----------------------------------------------------------------------===// + /// Options for the Sparsification pass. struct SparsificationOptions { SparsificationOptions(SparseParallelizationStrategy p, @@ -77,6 +86,19 @@ std::unique_ptr createSparsificationPass(const SparsificationOptions &options); +//===----------------------------------------------------------------------===// +// The PostSparsificationRewriting pass. +//===----------------------------------------------------------------------===// + +void populatePostSparsificationRewriting(RewritePatternSet &patterns, + bool enableRT, bool enableForeach, + bool enableConvert); + +std::unique_ptr createPostSparsificationRewritePass(); +std::unique_ptr +createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true, + bool enableConvert = true); + //===----------------------------------------------------------------------===// // The SparseTensorConversion pass. //===----------------------------------------------------------------------===// @@ -94,17 +116,6 @@ /// directly via the algorithm in ; /// however, beware that there are many formats not supported by this /// conversion method. -/// -/// The presence of the `kAuto` option violates our usual goal of keeping -/// policy completely separated from mechanism. The reason it exists is -/// because (at present) this strategy can only be specified on a per-file -/// basis. To see why this is a problem, note that `kDirect` cannot -/// support certain conversions; so if there is no `kAuto` setting, -/// then whenever a file contains a single non-`kDirect`-able conversion -/// the user would be forced to use `kViaCOO` for all conversions in -/// that file! In the future, instead of using this enum as a `Pass` -/// option, we could instead move it to being an attribute on the -/// conversion op; at which point `kAuto` would no longer be necessary. enum class SparseToSparseConversionStrategy { kAuto, kViaCOO, kDirect }; /// Converts command-line sparse2sparse flag to the strategy enum. @@ -140,7 +151,7 @@ SparseTensorTypeToBufferConverter(); }; -/// Sets up sparse tensor conversion rules. +/// Sets up sparse tensor codegen rules. void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter, RewritePatternSet &patterns, bool createSparseDeallocs, @@ -152,25 +163,42 @@ bool enableBufferInitialization); //===----------------------------------------------------------------------===// -// The PreSparsificationRewriting pass. +// The SparseBufferRewrite pass. //===----------------------------------------------------------------------===// -void populatePreSparsificationRewriting(RewritePatternSet &patterns); +void populateSparseBufferRewriting(RewritePatternSet &patterns, + bool enableBufferInitialization); -std::unique_ptr createPreSparsificationRewritePass(); +std::unique_ptr createSparseBufferRewritePass(); +std::unique_ptr +createSparseBufferRewritePass(bool enableBufferInitialization); //===----------------------------------------------------------------------===// -// The PostSparsificationRewriting pass. +// The SparseVectorization pass. //===----------------------------------------------------------------------===// -void populatePostSparsificationRewriting(RewritePatternSet &patterns, - bool enableRT, bool enableForeach, - bool enableConvert); +void populateSparseVectorizationPatterns(RewritePatternSet &patterns, + unsigned vectorLength, + bool enableVLAVectorization, + bool enableSIMDIndex32); -std::unique_ptr createPostSparsificationRewritePass(); -std::unique_ptr -createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true, - bool enableConvert = true); +std::unique_ptr createSparseVectorizationPass(); +std::unique_ptr createSparseVectorizationPass(unsigned vectorLength, + bool enableVLAVectorization, + bool enableSIMDIndex32); + +//===----------------------------------------------------------------------===// +// The SparseGPU pass. +//===----------------------------------------------------------------------===// + +void populateSparseGPUCodegenPatterns(RewritePatternSet &patterns, + unsigned numThreads); + +void populateSparseGPULibgenPatterns(RewritePatternSet &patterns, bool enableRT, + GPUDataTransferStrategy gpuDataTransfer); + +std::unique_ptr createSparseGPUCodegenPass(); +std::unique_ptr createSparseGPUCodegenPass(unsigned numThreads); //===----------------------------------------------------------------------===// // The SparseStorageSpecifierToLLVM pass. @@ -186,7 +214,11 @@ std::unique_ptr createStorageSpecifierToLLVMPass(); //===----------------------------------------------------------------------===// -// Other rewriting rules and passes. +// 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. //===----------------------------------------------------------------------===// std::unique_ptr createSparsificationAndBufferizationPass( @@ -197,32 +229,6 @@ bool enableBufferInitialization, unsigned vectorLength, bool enableVLAVectorization, bool enableSIMDIndex32); -void populateSparseBufferRewriting(RewritePatternSet &patterns, - bool enableBufferInitialization); - -std::unique_ptr createSparseBufferRewritePass(); -std::unique_ptr -createSparseBufferRewritePass(bool enableBufferInitialization); - -void populateSparseVectorizationPatterns(RewritePatternSet &patterns, - unsigned vectorLength, - bool enableVLAVectorization, - bool enableSIMDIndex32); - -std::unique_ptr createSparseVectorizationPass(); -std::unique_ptr createSparseVectorizationPass(unsigned vectorLength, - bool enableVLAVectorization, - bool enableSIMDIndex32); - -void populateSparseGPUCodegenPatterns(RewritePatternSet &patterns, - unsigned numThreads); - -void populateSparseGPULibgenPatterns(RewritePatternSet &patterns, bool enableRT, - GPUDataTransferStrategy gpuDataTransfer); - -std::unique_ptr createSparseGPUCodegenPass(); -std::unique_ptr createSparseGPUCodegenPass(unsigned numThreads); - //===----------------------------------------------------------------------===// // Registration. //===----------------------------------------------------------------------===// 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 @@ -123,7 +123,6 @@ ]; } - def PostSparsificationRewrite : Pass<"post-sparsification-rewrite", "ModuleOp"> { let summary = "Applies sparse tensor rewriting rules after sparsification"; let description = [{