[Mlir-commits] [mlir] [mlir][sparse] add sparsification options to pretty print and debug s… (PR #80205)

[Peiming Liu]

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

>From ddaa6073a0f2ebe79c3d838424d421699f8b536d Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming at google.com>
Date: Wed, 31 Jan 2024 18:32:01 +0000
Subject: [PATCH 1/2] [mlir][sparse] add sparsification options to pretty print
 and debug sparse iterations.

---
 .../Dialect/SparseTensor/Transforms/Passes.h  | 19 +++++++++++++++++--
 .../Dialect/SparseTensor/Transforms/Passes.td |  7 +++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index e93e2aefb344f..8b2875a751d4a 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -47,6 +47,12 @@ enum class ReinterpretMapScope {
   kExceptGeneric, // reinterprets operation other than linalg.generic
 };
 
+/// Defines a scope for reinterpret map pass.
+enum class DebugSparseIteration {
+  kNone,          // generate fully inlined (and functional) sparse iteration
+  kInterfaceOnly, // generate only place-holder for sparse iteration
+};
+
 #define GEN_PASS_DECL
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
@@ -74,11 +80,20 @@ std::unique_ptr<Pass> createPreSparsificationRewritePass();
 
 /// Options for the Sparsification pass.
 struct SparsificationOptions {
+  SparsificationOptions(SparseParallelizationStrategy p, DebugSparseIteration d,
+                        bool enableRT)
+      : parallelizationStrategy(p), debugSparseIteration(d),
+        enableRuntimeLibrary(enableRT) {}
+
   SparsificationOptions(SparseParallelizationStrategy p, bool enableRT)
-      : parallelizationStrategy(p), enableRuntimeLibrary(enableRT) {}
+      : SparsificationOptions(p, DebugSparseIteration::kNone, enableRT) {}
+
   SparsificationOptions()
-      : SparsificationOptions(SparseParallelizationStrategy::kNone, true) {}
+      : SparsificationOptions(SparseParallelizationStrategy::kNone,
+                              DebugSparseIteration::kNone, true) {}
+
   SparseParallelizationStrategy parallelizationStrategy;
+  DebugSparseIteration debugSparseIteration;
   bool enableRuntimeLibrary;
 };
 
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index f38779ed9ed2b..126b91510d391 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -130,6 +130,13 @@ def SparsificationPass : Pass<"sparsification", "ModuleOp"> {
              clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
                         "any-storage-any-loop",
                         "Enable sparse parallelization for any storage and loop."))}]>,
+    Option<"debugSparseIteration", "debug-sparse-iteration", "mlir::DebugSparseIteration",
+           "mlir::DebugSparseIteration::kNone",
+           "Pretty print sparse loops to debug sparse iteration", [{llvm::cl::values(
+             clEnumValN(mlir::DebugSparseIteration::kNone, "none",
+                        "Turn off pretty printing and generates functional code."),
+             clEnumValN(mlir::DebugSparseIteration::kInterfaceOnly, "interface-only",
+                        "Generate non-functional interfaces for sparse iteration."))}]>,
     Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
            "true", "Enable runtime library for manipulating sparse tensors">,
   ];

>From a457fc2ddedb5a5e44bc643bc4adf592746a3b0d Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming at google.com>
Date: Thu, 1 Feb 2024 19:14:33 +0000
Subject: [PATCH 2/2] address comments

---
 .../Dialect/SparseTensor/Transforms/Passes.h     | 16 ++++++++--------
 .../Dialect/SparseTensor/Transforms/Passes.td    | 14 +++++++-------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 8b2875a751d4a..e6c818d4abeac 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -48,9 +48,9 @@ enum class ReinterpretMapScope {
 };
 
 /// Defines a scope for reinterpret map pass.
-enum class DebugSparseIteration {
-  kNone,          // generate fully inlined (and functional) sparse iteration
-  kInterfaceOnly, // generate only place-holder for sparse iteration
+enum class SparseEmitStrategy {
+  kFunctional,     // generate fully inlined (and functional) sparse iteration
+  kDebugInterface, // generate only place-holder for sparse iteration
 };
 
 #define GEN_PASS_DECL
@@ -80,20 +80,20 @@ std::unique_ptr<Pass> createPreSparsificationRewritePass();
 
 /// Options for the Sparsification pass.
 struct SparsificationOptions {
-  SparsificationOptions(SparseParallelizationStrategy p, DebugSparseIteration d,
+  SparsificationOptions(SparseParallelizationStrategy p, SparseEmitStrategy d,
                         bool enableRT)
-      : parallelizationStrategy(p), debugSparseIteration(d),
+      : parallelizationStrategy(p), sparseEmitStrategy(d),
         enableRuntimeLibrary(enableRT) {}
 
   SparsificationOptions(SparseParallelizationStrategy p, bool enableRT)
-      : SparsificationOptions(p, DebugSparseIteration::kNone, enableRT) {}
+      : SparsificationOptions(p, SparseEmitStrategy::kFunctional, enableRT) {}
 
   SparsificationOptions()
       : SparsificationOptions(SparseParallelizationStrategy::kNone,
-                              DebugSparseIteration::kNone, true) {}
+                              SparseEmitStrategy::kFunctional, true) {}
 
   SparseParallelizationStrategy parallelizationStrategy;
-  DebugSparseIteration debugSparseIteration;
+  SparseEmitStrategy sparseEmitStrategy;
   bool enableRuntimeLibrary;
 };
 
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 126b91510d391..b7a14b89c2da2 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -130,13 +130,13 @@ def SparsificationPass : Pass<"sparsification", "ModuleOp"> {
              clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
                         "any-storage-any-loop",
                         "Enable sparse parallelization for any storage and loop."))}]>,
-    Option<"debugSparseIteration", "debug-sparse-iteration", "mlir::DebugSparseIteration",
-           "mlir::DebugSparseIteration::kNone",
-           "Pretty print sparse loops to debug sparse iteration", [{llvm::cl::values(
-             clEnumValN(mlir::DebugSparseIteration::kNone, "none",
-                        "Turn off pretty printing and generates functional code."),
-             clEnumValN(mlir::DebugSparseIteration::kInterfaceOnly, "interface-only",
-                        "Generate non-functional interfaces for sparse iteration."))}]>,
+    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."),
+             clEnumValN(mlir::SparseEmitStrategy::kDebugInterface, "debug-interface",
+                        "Emit non-functional but easy-to-read interfaces to debug."))}]>,
     Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
            "true", "Enable runtime library for manipulating sparse tensors">,
   ];