[Mlir-commits] [mlir] [mlir][sparse] add GPU num threads to sparsifier options (PR #189078)

[Vito Secona]

https://github.com/secona updated https://github.com/llvm/llvm-project/pull/189078

>From ae65703ec8e4f8adbc34d13738883a445ac4716d Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Thu, 19 Mar 2026 17:48:38 +0700
Subject: [PATCH 1/2] add numthreads to gpu codegen options

---
 mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h    | 5 +++++
 .../Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp | 3 ++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
index 2e76985e92e19..f17bc370648d6 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
@@ -151,6 +151,11 @@ struct SparsifierOptions : public PassPipelineOptions<SparsifierOptions> {
       desc("Enables GPU acceleration by means of direct library calls (like "
            "cuSPARSE)")};
 
+  /// This option is used to specify the number of threads of GPU codegen.
+  PassOptions::Option<unsigned> gpuNumThreads{
+      *this, "gpu-num-threads", desc("Number of threads for GPU codegen"),
+      init(1024)};
+
   /// Projects out the options for `createSparsificationPass`.
   SparsificationOptions sparsificationOptions() const {
     return SparsificationOptions(parallelization, emitStrategy,
diff --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
index dabbea1bdec63..d85966ec88f98 100644
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -53,7 +53,8 @@ void mlir::sparse_tensor::buildSparsifier(OpPassManager &pm,
   // GPU code generation.
   const bool gpuCodegen = options.gpuTriple.hasValue();
   if (gpuCodegen) {
-    pm.addPass(createSparseGPUCodegenPass());
+    pm.addPass(createSparseGPUCodegenPass(options.gpuNumThreads,
+                                          options.enableRuntimeLibrary));
     pm.addNestedPass<gpu::GPUModuleOp>(createStripDebugInfoPass());
     pm.addNestedPass<gpu::GPUModuleOp>(createSCFToControlFlowPass());
     pm.addNestedPass<gpu::GPUModuleOp>(createConvertGpuOpsToNVVMOps());

>From ce5b59abbe77736cac126bc8ec0e4ed7a1b7f543 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Tue, 31 Mar 2026 08:54:54 +0700
Subject: [PATCH 2/2] update docs for zero case

---
 mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
index f17bc370648d6..6bb46299738d5 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
@@ -153,7 +153,9 @@ struct SparsifierOptions : public PassPipelineOptions<SparsifierOptions> {
 
   /// This option is used to specify the number of threads of GPU codegen.
   PassOptions::Option<unsigned> gpuNumThreads{
-      *this, "gpu-num-threads", desc("Number of threads for GPU codegen"),
+      *this, "gpu-num-threads",
+      desc("Number of threads for GPU codegen. Setting this to 0 enables "
+           "direct library calls instead."),
       init(1024)};
 
   /// Projects out the options for `createSparsificationPass`.