[mlir][sparse] add parallelization options to mini pipeline (#104233)

Author: yinying-lisa-li

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

@@ -263,7 +263,8 @@ std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
     bool createSparseDeallocs, bool enableRuntimeLibrary,
     bool enableBufferInitialization, unsigned vectorLength,
     bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen,
-    SparseEmitStrategy emitStrategy);
+    SparseEmitStrategy emitStrategy,
+    SparseParallelizationStrategy parallelizationStrategy);
 
 //===----------------------------------------------------------------------===//
 // Sparse Iteration Transform Passes

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

@@ -496,6 +496,23 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
                         "Emit (experimental) loops (with sparse.iterate)."),
              clEnumValN(mlir::SparseEmitStrategy::kDebugInterface, "debug-interface",
                         "Emit non-functional but easy-to-read interfaces to debug."))}]>,
+    Option<"parallelization", "parallelization-strategy", "mlir::SparseParallelizationStrategy",
+           "mlir::SparseParallelizationStrategy::kNone",
+           "Set the parallelization strategy", [{llvm::cl::values(
+             clEnumValN(mlir::SparseParallelizationStrategy::kNone, "none",
+                        "Turn off sparse parallelization."),
+             clEnumValN(mlir::SparseParallelizationStrategy::kDenseOuterLoop,
+                        "dense-outer-loop",
+                        "Enable dense outer loop sparse parallelization."),
+             clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageOuterLoop,
+                        "any-storage-outer-loop",
+                        "Enable sparse parallelization regardless of storage for the outer loop."),
+             clEnumValN(mlir::SparseParallelizationStrategy::kDenseAnyLoop,
+                        "dense-any-loop",
+                        "Enable dense parallelization for any loop."),
+             clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
+                        "any-storage-any-loop",
+                        "Enable sparse parallelization for any storage and loop."))}]>,
   ];
 }
 

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

@@ -45,7 +45,8 @@ void mlir::sparse_tensor::buildSparsifier(OpPassManager &pm,
       /*enableVLAVectorization=*/options.armSVE,
       /*enableSIMDIndex32=*/options.force32BitVectorIndices,
       options.enableGPULibgen,
-      options.sparsificationOptions().sparseEmitStrategy));
+      options.sparsificationOptions().sparseEmitStrategy,
+      options.sparsificationOptions().parallelizationStrategy));
 
   // Bail-early for test setup.
   if (options.testBufferizationAnalysisOnly)

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

@@ -78,7 +78,8 @@ class SparsificationAndBufferizationPass
       const SparsificationOptions &sparsificationOptions,
       bool createSparseDeallocs, bool enableRuntimeLibrary,
       bool enableBufferInitialization, unsigned vl, bool vla, bool index32,
-      bool gpu, SparseEmitStrategy emitStrategy)
+      bool gpu, SparseEmitStrategy emitStrategy,
+      SparseParallelizationStrategy parallelizationStrategy)
       : bufferizationOptions(bufferizationOptions),
         sparsificationOptions(sparsificationOptions),
         createSparseDeallocs(createSparseDeallocs),
@@ -90,6 +91,7 @@ class SparsificationAndBufferizationPass
     enableSIMDIndex32 = index32;
     enableGPULibgen = gpu;
     sparseEmitStrategy = emitStrategy;
+    parallelization = parallelizationStrategy;
   }
 
   /// Bufferize all dense ops. This assumes that no further analysis is needed
@@ -124,6 +126,9 @@ class SparsificationAndBufferizationPass
     // Overrides the default emit strategy using user-provided value.
     this->sparsificationOptions.sparseEmitStrategy = sparseEmitStrategy;
 
+    // Overrides the default parallelization strategy using user-provided value.
+    this->sparsificationOptions.parallelizationStrategy = parallelization;
+
     // Run enabling transformations.
     {
       OpPassManager pm("builtin.module");
@@ -248,10 +253,12 @@ std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(
     bool createSparseDeallocs, bool enableRuntimeLibrary,
     bool enableBufferInitialization, unsigned vectorLength,
     bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen,
-    SparseEmitStrategy emitStrategy) {
+    SparseEmitStrategy emitStrategy,
+    SparseParallelizationStrategy parallelizationStrategy) {
   return std::make_unique<
       mlir::sparse_tensor::SparsificationAndBufferizationPass>(
       bufferizationOptions, sparsificationOptions, createSparseDeallocs,
       enableRuntimeLibrary, enableBufferInitialization, vectorLength,
-      enableVLAVectorization, enableSIMDIndex32, enableGPULibgen, emitStrategy);
+      enableVLAVectorization, enableSIMDIndex32, enableGPULibgen, emitStrategy,
+      parallelizationStrategy);
 }

mlir/test/Dialect/SparseTensor/minipipeline_parallel.mlir

@@ -0,0 +1,38 @@
+// RUN: mlir-opt %s --sparsification-and-bufferization        | FileCheck %s --check-prefix=CHECK-NOPARA
+// RUN: mlir-opt %s --sparsification-and-bufferization="parallelization-strategy=any-storage-any-loop" | FileCheck %s --check-prefix=CHECK-PARA
+
+// Test to ensure we can pass parallelization flags into
+// the mini sparsification and bufferization pipeline.
+
+#SparseMatrix = #sparse_tensor.encoding<{
+  map = (d0, d1) -> (d0 : compressed, d1 : compressed)
+}>
+
+#trait_ss = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A
+    affine_map<(i,j) -> (i,j)>   // X (out)
+  ],
+  iterator_types = ["parallel", "parallel"],
+  doc = "X(i,j) = A(i,j) * SCALE"
+}
+
+//
+// CHECK-NOPARA-LABEL: func.func @scale_ss
+// CHECK-NOPARA:       scf.for
+//
+// CHECK-PARA-LABEL: func.func @scale_ss
+// CHECK-PARA:       scf.parallel
+//
+func.func @scale_ss(%scale: f32,
+               %arga: tensor<?x?xf32, #SparseMatrix>,
+	       %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.generic #trait_ss
+     ins(%arga: tensor<?x?xf32, #SparseMatrix>)
+    outs(%argx: tensor<?x?xf32>) {
+      ^bb(%a: f32, %x: f32):
+        %0 = arith.mulf %a, %scale : f32
+        linalg.yield %0 : f32
+  } -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}