diff options
author | Peiming Liu <peiming@google.com> | 2024-06-20 14:23:36 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-06-20 14:23:36 -0700 |
commit | 1ba2768c638a3e0f98ccd002d3f328a18739f7a9 (patch) | |
tree | a64017284c4769c28672fa725a76c99f46b2766d /mlir | |
parent | fd001c16fed19249eec47d16ade9ad5810489c37 (diff) | |
download | llvm-1ba2768c638a3e0f98ccd002d3f328a18739f7a9.zip llvm-1ba2768c638a3e0f98ccd002d3f328a18739f7a9.tar.gz llvm-1ba2768c638a3e0f98ccd002d3f328a18739f7a9.tar.bz2 |
[mlir][sparse] expose emit strategy option to mini pipeline (#96238)
Diffstat (limited to 'mlir')
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 2edb9a6..8413691 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 7173cde..75a479b 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 ab4d90e..c5eb965 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 3ae6732..e088328 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp @@ -78,7 +78,7 @@ public: 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 @@ public: enableVLAVectorization = vla; enableSIMDIndex32 = index32; enableGPULibgen = gpu; + sparseEmitStrategy = emitStrategy; } /// Bufferize all dense ops. This assumes that no further analysis is needed @@ -120,6 +121,9 @@ public: } 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 77a0e89..6d8f9018 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) -> ( |