diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h index e93e2aefb344fd..e6c818d4abeacf 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 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" @@ -74,11 +80,20 @@ std::unique_ptr 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; }; diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td index f38779ed9ed2b8..b7a14b89c2da26 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<"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">, ];