[mlir][sparse] add parallelization options to mini pipeline

[yinying-lisa-li]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sparse

Author: Yinying Li (yinying-lisa-li)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/104233.diff

1 Files Affected:

• (modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td (+17)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 8ec18a1e186481..71617ee8a63d0d 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -487,6 +487,23 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
            "Enable i32 indexing into vectors (for efficient gather/scatter)">,
     Option<"enableGPULibgen", "enable-gpu-libgen", "bool", "false",
            "Enable GPU acceleration by means of direct library calls">,
+    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."))}]>,
     Option<"sparseEmitStrategy", "sparse-emit-strategy", "mlir::SparseEmitStrategy",
            "mlir::SparseEmitStrategy::kFunctional",
            "Emit functional code or interfaces (to debug) for sparse loops", [{llvm::cl::values(

[aartbik]

Please add a sanity test that the flags "take", see e.g. minipipeline_vector.mlir

[ingomueller-net]

Can I suggest to factor out the various parameters of createSparsificationAndBufferizationPass into a struct? This would make calls to that function forward compatible with any future parameters you might add.

[PeimingLiu]

Can I suggest to factor out the various parameters of createSparsificationAndBufferizationPass into a struct? This would make calls to that function forward compatible with any future parameters you might add.

SG! I will see whether I can find some cycles to do that!

[yinying-lisa-li]

Thank you both!