[Mlir-commits] [mlir] 8f0c014 - [mlir][sparse] add parallelization options to mini pipeline (#104233)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 30 14:47:40 PDT 2024
Author: Yinying Li
Date: 2024-08-30T17:47:37-04:00
New Revision: 8f0c014b12663129d8bfe0cc89f06e7a1d8b48c2
URL: https://github.com/llvm/llvm-project/commit/8f0c014b12663129d8bfe0cc89f06e7a1d8b48c2
DIFF: https://github.com/llvm/llvm-project/commit/8f0c014b12663129d8bfe0cc89f06e7a1d8b48c2.diff
LOG: [mlir][sparse] add parallelization options to mini pipeline (#104233)
Added:
mlir/test/Dialect/SparseTensor/minipipeline_parallel.mlir
Modified:
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 8413691910189a..d22df6a7857c1d 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -263,7 +263,8 @@ std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
bool createSparseDeallocs, bool enableRuntimeLibrary,
bool enableBufferInitialization, unsigned vectorLength,
bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen,
- SparseEmitStrategy emitStrategy);
+ SparseEmitStrategy emitStrategy,
+ SparseParallelizationStrategy parallelizationStrategy);
//===----------------------------------------------------------------------===//
// 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 8ec18a1e186481..a534381bd5c2f3 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -496,6 +496,23 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
"Emit (experimental) loops (with sparse.iterate)."),
clEnumValN(mlir::SparseEmitStrategy::kDebugInterface, "debug-interface",
"Emit non-functional but easy-to-read interfaces to debug."))}]>,
+ Option<"parallelization", "parallelization-strategy", "mlir::SparseParallelizationStrategy",
+ "mlir::SparseParallelizationStrategy::kNone",
+ "Set the parallelization strategy", [{llvm::cl::values(
+ clEnumValN(mlir::SparseParallelizationStrategy::kNone, "none",
+ "Turn off sparse parallelization."),
+ clEnumValN(mlir::SparseParallelizationStrategy::kDenseOuterLoop,
+ "dense-outer-loop",
+ "Enable dense outer loop sparse parallelization."),
+ clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageOuterLoop,
+ "any-storage-outer-loop",
+ "Enable sparse parallelization regardless of storage for the outer loop."),
+ clEnumValN(mlir::SparseParallelizationStrategy::kDenseAnyLoop,
+ "dense-any-loop",
+ "Enable dense parallelization for any loop."),
+ clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
+ "any-storage-any-loop",
+ "Enable sparse parallelization for any storage and loop."))}]>,
];
}
diff --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
index 12e330ac7efbdf..abc4a4c252841b 100644
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -45,7 +45,8 @@ void mlir::sparse_tensor::buildSparsifier(OpPassManager &pm,
/*enableVLAVectorization=*/options.armSVE,
/*enableSIMDIndex32=*/options.force32BitVectorIndices,
options.enableGPULibgen,
- options.sparsificationOptions().sparseEmitStrategy));
+ options.sparsificationOptions().sparseEmitStrategy,
+ options.sparsificationOptions().parallelizationStrategy));
// 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 e088328848c9c8..6e882a8d0ff30a 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -78,7 +78,8 @@ class SparsificationAndBufferizationPass
const SparsificationOptions &sparsificationOptions,
bool createSparseDeallocs, bool enableRuntimeLibrary,
bool enableBufferInitialization, unsigned vl, bool vla, bool index32,
- bool gpu, SparseEmitStrategy emitStrategy)
+ bool gpu, SparseEmitStrategy emitStrategy,
+ SparseParallelizationStrategy parallelizationStrategy)
: bufferizationOptions(bufferizationOptions),
sparsificationOptions(sparsificationOptions),
createSparseDeallocs(createSparseDeallocs),
@@ -90,6 +91,7 @@ class SparsificationAndBufferizationPass
enableSIMDIndex32 = index32;
enableGPULibgen = gpu;
sparseEmitStrategy = emitStrategy;
+ parallelization = parallelizationStrategy;
}
/// Bufferize all dense ops. This assumes that no further analysis is needed
@@ -124,6 +126,9 @@ class SparsificationAndBufferizationPass
// Overrides the default emit strategy using user-provided value.
this->sparsificationOptions.sparseEmitStrategy = sparseEmitStrategy;
+ // Overrides the default parallelization strategy using user-provided value.
+ this->sparsificationOptions.parallelizationStrategy = parallelization;
+
// Run enabling transformations.
{
OpPassManager pm("builtin.module");
@@ -248,10 +253,12 @@ std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(
bool createSparseDeallocs, bool enableRuntimeLibrary,
bool enableBufferInitialization, unsigned vectorLength,
bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen,
- SparseEmitStrategy emitStrategy) {
+ SparseEmitStrategy emitStrategy,
+ SparseParallelizationStrategy parallelizationStrategy) {
return std::make_unique<
mlir::sparse_tensor::SparsificationAndBufferizationPass>(
bufferizationOptions, sparsificationOptions, createSparseDeallocs,
enableRuntimeLibrary, enableBufferInitialization, vectorLength,
- enableVLAVectorization, enableSIMDIndex32, enableGPULibgen, emitStrategy);
+ enableVLAVectorization, enableSIMDIndex32, enableGPULibgen, emitStrategy,
+ parallelizationStrategy);
}
diff --git a/mlir/test/Dialect/SparseTensor/minipipeline_parallel.mlir b/mlir/test/Dialect/SparseTensor/minipipeline_parallel.mlir
new file mode 100644
index 00000000000000..d97d6e58a3df2d
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/minipipeline_parallel.mlir
@@ -0,0 +1,38 @@
+// RUN: mlir-opt %s --sparsification-and-bufferization | FileCheck %s --check-prefix=CHECK-NOPARA
+// RUN: mlir-opt %s --sparsification-and-bufferization="parallelization-strategy=any-storage-any-loop" | FileCheck %s --check-prefix=CHECK-PARA
+
+// Test to ensure we can pass parallelization flags into
+// the mini sparsification and bufferization pipeline.
+
+#SparseMatrix = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed, d1 : compressed)
+}>
+
+#trait_ss = {
+ indexing_maps = [
+ affine_map<(i,j) -> (i,j)>, // A
+ affine_map<(i,j) -> (i,j)> // X (out)
+ ],
+ iterator_types = ["parallel", "parallel"],
+ doc = "X(i,j) = A(i,j) * SCALE"
+}
+
+//
+// CHECK-NOPARA-LABEL: func.func @scale_ss
+// CHECK-NOPARA: scf.for
+//
+// CHECK-PARA-LABEL: func.func @scale_ss
+// CHECK-PARA: scf.parallel
+//
+func.func @scale_ss(%scale: f32,
+ %arga: tensor<?x?xf32, #SparseMatrix>,
+ %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.generic #trait_ss
+ ins(%arga: tensor<?x?xf32, #SparseMatrix>)
+ outs(%argx: tensor<?x?xf32>) {
+ ^bb(%a: f32, %x: f32):
+ %0 = arith.mulf %a, %scale : f32
+ linalg.yield %0 : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
More information about the Mlir-commits
mailing list