[mlir][sparse] cleanup of pass code

[aartbik]

Summary

This reorders the pass declarations and definitions in a
consistent order and makes the layout a bit more clear.
Also makes the "minipipeline" sparsification-and-bufferization
available through the command line for individual testing.

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

 //===- Passes.h - Sparse tensor pass entry points ---------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // 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_
 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
 
 #include "mlir/IR/PatternMatch.h"
 #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
 /// outermost loop only), and (2) the generated code is an actual for-loop
 /// (and not a co-iterating while-loop).
 enum class SparseParallelizationStrategy {
   kNone,
   kDenseOuterLoop,
   kAnyStorageOuterLoop,
   kDenseAnyLoop,
   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<Pass> createPreSparsificationRewritePass();
+
+//===----------------------------------------------------------------------===//
+// The Sparsification pass.
+//===----------------------------------------------------------------------===//
+
 /// Options for the Sparsification pass.
 struct SparsificationOptions {
   SparsificationOptions(SparseParallelizationStrategy p,
                         GPUDataTransferStrategy t, bool idxReduc,
                         bool gpuLibgen, bool enableRT)
       : parallelizationStrategy(p), gpuDataTransferStrategy(t),
         enableIndexReduction(idxReduc), enableGPULibgen(gpuLibgen),
         enableRuntimeLibrary(enableRT) {}
void populateSparsificationPatterns(
     RewritePatternSet &patterns,
     const SparsificationOptions &options = SparsificationOptions());
 
 std::unique_ptr<Pass> createSparsificationPass();
 std::unique_ptr<Pass>
 createSparsificationPass(const SparsificationOptions &options);
 
 //===----------------------------------------------------------------------===//
+// The PostSparsificationRewriting pass.
+//===----------------------------------------------------------------------===//
+
+void populatePostSparsificationRewriting(RewritePatternSet &patterns,
+                                         bool enableRT, bool enableForeach,
+                                         bool enableConvert);
+
+std::unique_ptr<Pass> createPostSparsificationRewritePass();
+std::unique_ptr<Pass>
+createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true,
+                                    bool enableConvert = true);
+
+//===----------------------------------------------------------------------===//
 // The SparseTensorConversion pass.
 //===----------------------------------------------------------------------===//
 
 /// Sparse tensor type converter into an opaque pointer.
 class SparseTensorTypeToPtrConverter : public TypeConverter {
 public:
   SparseTensorTypeToPtrConverter();
 };
 
 /// Defines a strategy for implementing sparse-to-sparse conversion.
 /// `kAuto` leaves it up to the compiler to automatically determine
 /// the method used.  `kViaCOO` converts the source tensor to COO and
 /// then converts the COO to the target format.  `kDirect` converts
 /// directly via the algorithm in <https://arxiv.org/abs/2001.02609>;
 /// 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.
 SparseToSparseConversionStrategy sparseToSparseConversionStrategy(int32_t flag);
 
 /// SparseTensorConversion options.
 struct SparseTensorConversionOptions {
   SparseTensorConversionOptions(SparseToSparseConversionStrategy s2s)

 //===----------------------------------------------------------------------===//
 
 /// Sparse tensor type converter into an actual buffer.
 class SparseTensorTypeToBufferConverter : public TypeConverter {
 public:
   SparseTensorTypeToBufferConverter();
 };
 
-/// Sets up sparse tensor conversion rules.
+/// Sets up sparse tensor codegen rules.
 void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                          RewritePatternSet &patterns,
                                          bool createSparseDeallocs,
                                          bool enableBufferInitialization);
 
 std::unique_ptr<Pass> createSparseTensorCodegenPass();
 std::unique_ptr<Pass>
 createSparseTensorCodegenPass(bool createSparseDeallocs,
                               bool enableBufferInitialization);
 
 //===----------------------------------------------------------------------===//
-// The PreSparsificationRewriting pass.
+// The SparseBufferRewrite pass.
 //===----------------------------------------------------------------------===//
 
-void populatePreSparsificationRewriting(RewritePatternSet &patterns);
+void populateSparseBufferRewriting(RewritePatternSet &patterns,
+                                   bool enableBufferInitialization);
 
-std::unique_ptr<Pass> createPreSparsificationRewritePass();
+std::unique_ptr<Pass> createSparseBufferRewritePass();
+std::unique_ptr<Pass>
+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<Pass> createPostSparsificationRewritePass();
-std::unique_ptr<Pass>
-createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true,
-                                    bool enableConvert = true);
+std::unique_ptr<Pass> createSparseVectorizationPass();
+std::unique_ptr<Pass> 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<Pass> createSparseGPUCodegenPass();
+std::unique_ptr<Pass> createSparseGPUCodegenPass(unsigned numThreads);
 
 //===----------------------------------------------------------------------===//
 // The SparseStorageSpecifierToLLVM pass.
 //===----------------------------------------------------------------------===//
 
 class StorageSpecifierToLLVMTypeConverter : public TypeConverter {
 public:
   StorageSpecifierToLLVMTypeConverter();
 };
 
 void populateStorageSpecifierToLLVMPatterns(TypeConverter &converter,
                                             RewritePatternSet &patterns);
 std::unique_ptr<Pass> 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<Pass> createSparsificationAndBufferizationPass(
     const bufferization::OneShotBufferizationOptions &bufferizationOptions,
     const SparsificationOptions &sparsificationOptions,
     const SparseTensorConversionOptions &sparseTensorConversionOptions,
     bool createSparseDeallocs, bool enableRuntimeLibrary,
     bool enableBufferInitialization, unsigned vectorLength,
     bool enableVLAVectorization, bool enableSIMDIndex32);
 
-void populateSparseBufferRewriting(RewritePatternSet &patterns,
-                                   bool enableBufferInitialization);
-
-std::unique_ptr<Pass> createSparseBufferRewritePass();
-std::unique_ptr<Pass>
-createSparseBufferRewritePass(bool enableBufferInitialization);
-
-void populateSparseVectorizationPatterns(RewritePatternSet &patterns,
-                                         unsigned vectorLength,
-                                         bool enableVLAVectorization,
-                                         bool enableSIMDIndex32);
-
-std::unique_ptr<Pass> createSparseVectorizationPass();
-std::unique_ptr<Pass> createSparseVectorizationPass(unsigned vectorLength,
-                                                    bool enableVLAVectorization,
-                                                    bool enableSIMDIndex32);
-
-void populateSparseGPUCodegenPatterns(RewritePatternSet &patterns,
-                                      unsigned numThreads);
-
-void populateSparseGPULibgenPatterns(RewritePatternSet &patterns, bool enableRT,
-                                     GPUDataTransferStrategy gpuDataTransfer);
-
-std::unique_ptr<Pass> createSparseGPUCodegenPass();
-std::unique_ptr<Pass> createSparseGPUCodegenPass(unsigned numThreads);
-
 //===----------------------------------------------------------------------===//
 // Registration.
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

let options = [
     Option<"enableGPULibgen", "enable-gpu-libgen", "bool",
            "false",
            "Enable GPU acceleration by means of direct library calls (like cuSPARSE)">,
     Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
            "true", "Enable runtime library for manipulating sparse tensors">,
   ];
 }
 
-
 def PostSparsificationRewrite : Pass<"post-sparsification-rewrite", "ModuleOp"> {
   let summary = "Applies sparse tensor rewriting rules after sparsification";
   let description = [{
     A pass that applies rewriting rules to sparse tensor operations after
     running the actual sparsification pass.
   }];
   let constructor = "mlir::createPostSparsificationRewritePass()";
   let dependentDialects = [