Skip to content

Commit 20a9fa3

Browse files
authored
[mlir][sparse] add sparsification options to pretty print and debug s… (#80205)
…parse loops.
1 parent e296ced commit 20a9fa3

File tree

2 files changed

+24
-2
lines changed

2 files changed

+24
-2
lines changed

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

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,12 @@ enum class ReinterpretMapScope {
4747
kExceptGeneric, // reinterprets operation other than linalg.generic
4848
};
4949

50+
/// Defines a scope for reinterpret map pass.
51+
enum class SparseEmitStrategy {
52+
kFunctional, // generate fully inlined (and functional) sparse iteration
53+
kDebugInterface, // generate only place-holder for sparse iteration
54+
};
55+
5056
#define GEN_PASS_DECL
5157
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
5258

@@ -74,11 +80,20 @@ std::unique_ptr<Pass> createPreSparsificationRewritePass();
7480

7581
/// Options for the Sparsification pass.
7682
struct SparsificationOptions {
83+
SparsificationOptions(SparseParallelizationStrategy p, SparseEmitStrategy d,
84+
bool enableRT)
85+
: parallelizationStrategy(p), sparseEmitStrategy(d),
86+
enableRuntimeLibrary(enableRT) {}
87+
7788
SparsificationOptions(SparseParallelizationStrategy p, bool enableRT)
78-
: parallelizationStrategy(p), enableRuntimeLibrary(enableRT) {}
89+
: SparsificationOptions(p, SparseEmitStrategy::kFunctional, enableRT) {}
90+
7991
SparsificationOptions()
80-
: SparsificationOptions(SparseParallelizationStrategy::kNone, true) {}
92+
: SparsificationOptions(SparseParallelizationStrategy::kNone,
93+
SparseEmitStrategy::kFunctional, true) {}
94+
8195
SparseParallelizationStrategy parallelizationStrategy;
96+
SparseEmitStrategy sparseEmitStrategy;
8297
bool enableRuntimeLibrary;
8398
};
8499

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,13 @@ def SparsificationPass : Pass<"sparsification", "ModuleOp"> {
130130
clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
131131
"any-storage-any-loop",
132132
"Enable sparse parallelization for any storage and loop."))}]>,
133+
Option<"sparseEmitStrategy", "sparse-emit-strategy", "mlir::SparseEmitStrategy",
134+
"mlir::SparseEmitStrategy::kFunctional",
135+
"Emit functional code or interfaces (to debug) for sparse loops", [{llvm::cl::values(
136+
clEnumValN(mlir::SparseEmitStrategy::kFunctional, "functional",
137+
"Emit functional code."),
138+
clEnumValN(mlir::SparseEmitStrategy::kDebugInterface, "debug-interface",
139+
"Emit non-functional but easy-to-read interfaces to debug."))}]>,
133140
Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
134141
"true", "Enable runtime library for manipulating sparse tensors">,
135142
];

0 commit comments

Comments
 (0)