[mlir][sparse] Add new option (enable-runtime-library) to sparse compiler pipeline

[Peiming]

Summary

Add new option (enable-runtime-library) to sparse compiler pipeline, it allows us to decide whether we need to rewrite operations (e.g., concatenate, reshape) within sparsification (when using codegen) or convert them after sparsification (when using runtime library).

[Peiming]

Take the codegen path when RT is disabled.

[aartbik]

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

80

I love your enthusiasm adding this TODO ;-)

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

73

if we document the value (good practice), let's also document the other two false

and I would say, "enable" instead of "use" to stay in the same terminology as used in the flag description

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

102

Enable

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

55

move this comment and decl down too

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

308

If not enableRT (flag name) ....
or
If RT not enabled ....

[Peiming]

Address comments from Aart

[Peiming]

First codegen integrate test!

[Peiming]

add missing period...

[Peiming]

update test

[aartbik]

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

74

oops, supernit, we usually put it before the value, as in

/*disable SIMD Index32 = */false

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

68 ↗

(On Diff #459184)

can we actually run this always, even for lib cases to clean up?

This may impact some CHECK tests though, so I am okay either way

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_codegen_dim.mlir

14 ↗

(On Diff #459184)

Move down one line

also, rename MAT into DCSR

[Peiming]

Address comments from Aart.

[Peiming]

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

68 ↗

(On Diff #459184)

Actually it does not impact any CHECT tests, probably only end-to-end integrate test are involved here, so it wouldn't change the result.

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

struct SparseCompilerOptions
   PassOptions::Option<int32_t> vectorLength{
       *this, "vl", desc("Set the vector length"), init(1)};
   PassOptions::Option<bool> enableSIMDIndex32{
       *this, "enable-simd-index32",
       desc("Enable i32 indexing into vectors (for efficiency)"), init(false)};
   PassOptions::Option<bool> enableVLAVectorization{
       *this, "enable-vla-vectorization",
       desc("Enable vector length agnostic vectorization"), init(false)};
+  PassOptions::Option<bool> enableRuntimeLibrary{
+      *this, "enable-runtime-library",
+      // TODO: Disable runtime library by default after feature complete.
+      desc("Use runtime library for manipulating sparse tensors"), init(true)};
+
   PassOptions::Option<bool> testBufferizationAnalysisOnly{
       *this, "test-bufferization-analysis-only",
       desc("Run only the inplacability analysis"), init(false)};
 
   /// Projects out the options for `createSparsificationPass`.
   SparsificationOptions sparsificationOptions() const {
     return SparsificationOptions(parallelization, vectorization, vectorLength,
-                                 enableSIMDIndex32, enableVLAVectorization);
+                                 enableSIMDIndex32, enableVLAVectorization,
+                                 enableRuntimeLibrary);
   }
 
   // These options must be kept in sync with `SparseTensorConversionBase`.
   PassOptions::Option<int32_t> sparseToSparse{
       *this, "s2s-strategy",
       desc("Set the strategy for sparse-to-sparse conversion"), init(0)};
 
   /// Projects out the options for `createSparsificationPass`.

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

 #define GEN_PASS_DECL_SPARSETENSORCONVERSIONPASS
 #define GEN_PASS_DECL_SPARSETENSORCODEGEN
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
 /// Options for the Sparsification pass.
 struct SparsificationOptions {
   SparsificationOptions(SparseParallelizationStrategy p,
                         SparseVectorizationStrategy v, unsigned vl, bool e,
-                        bool vla)
+                        bool vla, bool rt)
       : parallelizationStrategy(p), vectorizationStrategy(v), vectorLength(vl),
-        enableSIMDIndex32(e), enableVLAVectorization(vla) {}
+        enableSIMDIndex32(e), enableVLAVectorization(vla),
+        enableRuntimeLibrary(rt) {}
   SparsificationOptions()
       : SparsificationOptions(SparseParallelizationStrategy::kNone,
                               SparseVectorizationStrategy::kNone, 1u, false,
-                              false) {}
+                              false, true /*use runtime library*/) {}
   SparseParallelizationStrategy parallelizationStrategy;
   SparseVectorizationStrategy vectorizationStrategy;
   unsigned vectorLength;
   bool enableSIMDIndex32;
   bool enableVLAVectorization;
+  bool enableRuntimeLibrary;
 };
 
 /// Sets up sparsification rewriting rules with the given options.
 void populateSparsificationPatterns(
     RewritePatternSet &patterns,
     const SparsificationOptions &options = SparsificationOptions());
 
 std::unique_ptr<Pass> createSparsificationPass();
void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                          RewritePatternSet &patterns);
 
 std::unique_ptr<Pass> createSparseTensorCodegenPass();
 
 //===----------------------------------------------------------------------===//
 // Other rewriting rules and passes.
 //===----------------------------------------------------------------------===//
 
-void populateSparseTensorRewriting(RewritePatternSet &patterns);
+void populateSparseTensorRewriting(RewritePatternSet &patterns, bool enableRT);
 
 std::unique_ptr<Pass> createDenseBufferizationPass(
     const bufferization::OneShotBufferizationOptions &options);
 
 //===----------------------------------------------------------------------===//
 // 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

Option<"vectorization", "vectorization-strategy", "mlir::SparseVectorizationStrategy",
              clEnumValN(mlir::SparseVectorizationStrategy::kAnyStorageInnerLoop,
                         "any-storage-inner-loop",
                         "Enable sparse vectorization for inner loops with any storage."))}]>,
     Option<"vectorLength", "vl", "int32_t", "1",
            "Set the vector length">,
     Option<"enableSIMDIndex32", "enable-simd-index32", "bool", "false",
            "Enable i32 indexing into vectors (for efficiency)">,
     Option<"enableVLAVectorization", "enable-vla-vectorization", "bool",
-           "false", "Enable vector length agnostic vectorization">
+           "false", "Enable vector length agnostic vectorization">,
+    Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
+           "true", "Use runtime library for manipulating sparse tensors">
   ];
 }
 
 def SparseTensorConversionPass : Pass<"sparse-tensor-conversion", "ModuleOp"> {
   let summary = "Convert sparse tensors and primitives to library calls";
   let description = [{
     A pass that converts sparse tensor primitives into calls into a runtime
     support library. Sparse tensor types are converted into opaque pointers

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

struct SparsificationPass
   SparsificationPass() = default;
   SparsificationPass(const SparsificationPass &pass) = default;
   SparsificationPass(const SparsificationOptions &options) {
     parallelization = options.parallelizationStrategy;
     vectorization = options.vectorizationStrategy;
     vectorLength = options.vectorLength;
     enableSIMDIndex32 = options.enableSIMDIndex32;
     enableVLAVectorization = options.enableVLAVectorization;
+    enableRuntimeLibrary = options.enableRuntimeLibrary;
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     // Apply pre-rewriting.
     RewritePatternSet prePatterns(ctx);
-    populateSparseTensorRewriting(prePatterns);
-    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(prePatterns));
     // Translate strategy flags to strategy options.
     SparsificationOptions options(parallelization, vectorization, vectorLength,
-                                  enableSIMDIndex32, enableVLAVectorization);
+                                  enableSIMDIndex32, enableVLAVectorization,
+                                  enableRuntimeLibrary);
+    populateSparseTensorRewriting(prePatterns, options.enableRuntimeLibrary);
+    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(prePatterns));
     // Apply sparsification and vector cleanup rewriting.
     RewritePatternSet patterns(ctx);
     populateSparsificationPatterns(patterns, options);
     vector::populateVectorToVectorCanonicalizationPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

 };
 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
 
-void mlir::populateSparseTensorRewriting(RewritePatternSet &patterns) {
+void mlir::populateSparseTensorRewriting(RewritePatternSet &patterns,
+                                         bool /*enableRT*/) {
   patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd,
                ReshapeRewriter<tensor::ExpandShapeOp>,
                ReshapeRewriter<tensor::CollapseShapeOp>>(patterns.getContext());
+  // If not enable RT, rewrite concatenate ops, etc here.
 }