[Mlir-commits] [mlir] [mlir][sparse] expose emit strategy option to mini pipeline (PR #96238)

[Peiming Liu]

https://github.com/PeimingLiu updated https://github.com/llvm/llvm-project/pull/96238

>From 2759e29964fbfdd45bd0fff1321a3adcef9eaa6b Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming at google.com>
Date: Thu, 20 Jun 2024 20:57:53 +0000
Subject: [PATCH] [mlir][sparse] expose emit strategy option to mini pipeline

---
 .../mlir/Dialect/SparseTensor/Transforms/Passes.h     |  3 ++-
 .../mlir/Dialect/SparseTensor/Transforms/Passes.td    |  9 +++++++++
 .../SparseTensor/Pipelines/SparseTensorPipelines.cpp  |  3 ++-
 .../Transforms/SparsificationAndBufferizationPass.cpp | 11 ++++++++---
 .../Dialect/SparseTensor/sparse_iteration_to_scf.mlir |  1 +
 5 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 2edb9a61d1876..8413691910189 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -262,7 +262,8 @@ std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
     const SparsificationOptions &sparsificationOptions,
     bool createSparseDeallocs, bool enableRuntimeLibrary,
     bool enableBufferInitialization, unsigned vectorLength,
-    bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen);
+    bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen,
+    SparseEmitStrategy emitStrategy);
 
 //===----------------------------------------------------------------------===//
 // Sparse Iteration Transform Passes
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 7173cde67384b..75a479b898040 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -475,6 +475,15 @@ 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<"sparseEmitStrategy", "sparse-emit-strategy", "mlir::SparseEmitStrategy",
+           "mlir::SparseEmitStrategy::kFunctional",
+           "Emit functional code or interfaces (to debug) for sparse loops", [{llvm::cl::values(
+             clEnumValN(mlir::SparseEmitStrategy::kFunctional, "functional",
+                        "Emit functional code (with scf.for/while)."),
+             clEnumValN(mlir::SparseEmitStrategy::kSparseIterator, "sparse-iterator",
+                        "Emit (experimental) loops (with sparse.iterate)."),
+             clEnumValN(mlir::SparseEmitStrategy::kDebugInterface, "debug-interface",
+                        "Emit non-functional but easy-to-read interfaces to debug."))}]>,
   ];
 }
 
diff --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
index ab4d90ec745d4..c5eb965884396 100644
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -44,7 +44,8 @@ void mlir::sparse_tensor::buildSparsifier(OpPassManager &pm,
       options.vectorLength,
       /*enableVLAVectorization=*/options.armSVE,
       /*enableSIMDIndex32=*/options.force32BitVectorIndices,
-      options.enableGPULibgen));
+      options.enableGPULibgen,
+      options.sparsificationOptions().sparseEmitStrategy));
 
   // Bail-early for test setup.
   if (options.testBufferizationAnalysisOnly)
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index 3ae6732f900fe..e088328848c9c 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -78,7 +78,7 @@ class SparsificationAndBufferizationPass
       const SparsificationOptions &sparsificationOptions,
       bool createSparseDeallocs, bool enableRuntimeLibrary,
       bool enableBufferInitialization, unsigned vl, bool vla, bool index32,
-      bool gpu)
+      bool gpu, SparseEmitStrategy emitStrategy)
       : bufferizationOptions(bufferizationOptions),
         sparsificationOptions(sparsificationOptions),
         createSparseDeallocs(createSparseDeallocs),
@@ -89,6 +89,7 @@ class SparsificationAndBufferizationPass
     enableVLAVectorization = vla;
     enableSIMDIndex32 = index32;
     enableGPULibgen = gpu;
+    sparseEmitStrategy = emitStrategy;
   }
 
   /// Bufferize all dense ops. This assumes that no further analysis is needed
@@ -120,6 +121,9 @@ class SparsificationAndBufferizationPass
   }
 
   void runOnOperation() override {
+    // Overrides the default emit strategy using user-provided value.
+    this->sparsificationOptions.sparseEmitStrategy = sparseEmitStrategy;
+
     // Run enabling transformations.
     {
       OpPassManager pm("builtin.module");
@@ -243,10 +247,11 @@ std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(
     const SparsificationOptions &sparsificationOptions,
     bool createSparseDeallocs, bool enableRuntimeLibrary,
     bool enableBufferInitialization, unsigned vectorLength,
-    bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen) {
+    bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen,
+    SparseEmitStrategy emitStrategy) {
   return std::make_unique<
       mlir::sparse_tensor::SparsificationAndBufferizationPass>(
       bufferizationOptions, sparsificationOptions, createSparseDeallocs,
       enableRuntimeLibrary, enableBufferInitialization, vectorLength,
-      enableVLAVectorization, enableSIMDIndex32, enableGPULibgen);
+      enableVLAVectorization, enableSIMDIndex32, enableGPULibgen, emitStrategy);
 }
diff --git a/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir b/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir
index 77a0e89dc7c81..6d8f9018ad3a5 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir
@@ -1,5 +1,6 @@
 // RUN: mlir-opt %s --lower-sparse-iteration-to-scf | FileCheck %s
 // RUN: mlir-opt %s --sparse-space-collapse --lower-sparse-iteration-to-scf | FileCheck %s --check-prefix COLLAPSED
+// RUN: mlir-opt %s --sparsification-and-bufferization="sparse-emit-strategy=sparse-iterator" | FileCheck %s --check-prefix COLLAPSED
 
 #COO = #sparse_tensor.encoding<{
   map = (i, j) -> (