Skip to content

Commit

Permalink
[mlir][sparse] add sparsification options to pretty print and debug s… (
Browse files Browse the repository at this point in the history
llvm#80205)

…parse loops.
  • Loading branch information
PeimingLiu authored and agozillon committed Feb 5, 2024
1 parent 47d67de commit 60c0c67
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 2 deletions.
19 changes: 17 additions & 2 deletions mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,12 @@ enum class ReinterpretMapScope {
kExceptGeneric, // reinterprets operation other than linalg.generic
};

/// Defines a scope for reinterpret map pass.
enum class SparseEmitStrategy {
kFunctional, // generate fully inlined (and functional) sparse iteration
kDebugInterface, // generate only place-holder for sparse iteration
};

#define GEN_PASS_DECL
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"

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

/// Options for the Sparsification pass.
struct SparsificationOptions {
SparsificationOptions(SparseParallelizationStrategy p, SparseEmitStrategy d,
bool enableRT)
: parallelizationStrategy(p), sparseEmitStrategy(d),
enableRuntimeLibrary(enableRT) {}

SparsificationOptions(SparseParallelizationStrategy p, bool enableRT)
: parallelizationStrategy(p), enableRuntimeLibrary(enableRT) {}
: SparsificationOptions(p, SparseEmitStrategy::kFunctional, enableRT) {}

SparsificationOptions()
: SparsificationOptions(SparseParallelizationStrategy::kNone, true) {}
: SparsificationOptions(SparseParallelizationStrategy::kNone,
SparseEmitStrategy::kFunctional, true) {}

SparseParallelizationStrategy parallelizationStrategy;
SparseEmitStrategy sparseEmitStrategy;
bool enableRuntimeLibrary;
};

Expand Down
7 changes: 7 additions & 0 deletions mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
Original file line number Diff line number Diff line change
Expand Up @@ -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<"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">,
];
Expand Down

0 comments on commit 60c0c67

Please sign in to comment.