[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

72–75

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

54–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

73

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

/*disable SIMD Index32 = */false

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

68

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

Move down one line

also, rename MAT into DCSR

[Peiming]

Address comments from Aart.

[Peiming]

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

68

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",
+      desc("Enable runtime library for manipulating sparse tensors"),
+      // TODO: Disable runtime library by default after feature complete.
+      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) {}
+                              SparseVectorizationStrategy::kNone, 1u,
+                              false, /*disable SIMD Index32*/
+                              false, /*disable VLA Vectorization*/
+                              true /*enable 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", "Enable 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/Pipelines/SparseTensorPipelines.cpp

 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Passes.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::sparse_tensor;
 
 /// Return configuration options for One-Shot Bufferize.
 static bufferization::OneShotBufferizationOptions
 getBufferizationOptions(bool analysisOnly) {
   using namespace bufferization;
void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   pm.addNestedPass<func::FuncOp>(createLinalgGeneralizationPass());
   pm.addPass(
       bufferization::createTensorCopyInsertionPass(getBufferizationOptions(
           /*analysisOnly=*/options.testBufferizationAnalysisOnly)));
   if (options.testBufferizationAnalysisOnly)
     return;
   pm.addPass(createSparsificationPass(options.sparsificationOptions()));
+  if (options.enableRuntimeLibrary)
     pm.addPass(createSparseTensorConversionPass(
         options.sparseTensorConversionOptions()));
+  else {
+    pm.addPass(createSparseTensorCodegenPass());
+    // To remove useless cast.
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  }
+
   pm.addPass(createDenseBufferizationPass(
       getBufferizationOptions(/*analysisOnly=*/false)));
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());
   // TODO(springerm): Add sparse support to the BufferDeallocation pass and add
   // it to this pipeline.
   pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());
   pm.addNestedPass<func::FuncOp>(createConvertVectorToSCFPass());

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);
+    // Apply pre-rewriting.
+    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());
+  // TODO: If RT not enabled, rewrite concatenate ops, etc here.
 }

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

+// Test with/without runtime library, the result should always be identical.
+
+// RUN: mlir-opt %s --sparse-compiler=enable-runtime-library=false | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+// RUN: mlir-opt %s --sparse-compiler=enable-runtime-library=true | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+#MAT  = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed"]
+}>
+
+module {
+  //
+  // Main driver.
+  //
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c2 = arith.constant 2 : index
+    %c3 = arith.constant 3 : index
+    %t1 = bufferization.alloc_tensor() : tensor<4x5xf64, #MAT>
+    %t2 = bufferization.alloc_tensor(%c2, %c3) : tensor<?x?xf64, #MAT>
+
+    %d1_0 = tensor.dim %t1, %c0 : tensor<4x5xf64, #MAT>
+    %d2_0 = tensor.dim %t2, %c0 : tensor<?x?xf64, #MAT>
+    %d1_1 = tensor.dim %t1, %c1 : tensor<4x5xf64, #MAT>
+    %d2_1 = tensor.dim %t2, %c1 : tensor<?x?xf64, #MAT>
+    // CHECK: 4
+    vector.print %d1_0 : index
+    // CHECK-NEXT: 2
+    vector.print %d2_0 : index
+    // CHECK-NEXT: 5
+    vector.print %d1_1 : index
+    // CHECK-NEXT: 3
+    vector.print %d2_1 : index
+    return
+  }
+}
+