[Mlir-commits] [mlir] [mlir][sparse] add parallelization options to mini pipeline (PR #104233)
Yinying Li
llvmlistbot at llvm.org
Wed Aug 14 17:21:54 PDT 2024
https://github.com/yinying-lisa-li updated https://github.com/llvm/llvm-project/pull/104233
>From 9c5056acd7f22bce0fdab4c2546c4569cd998d0e Mon Sep 17 00:00:00 2001
From: Yinying Li <yinyingli at google.com>
Date: Wed, 14 Aug 2024 20:46:03 +0000
Subject: [PATCH 1/2] add parallelization options
---
.../Dialect/SparseTensor/Transforms/Passes.td | 17 +++++++++++++++++
1 file changed, 17 insertions(+)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 8ec18a1e186481..71617ee8a63d0d 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -487,6 +487,23 @@ 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<"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."))}]>,
Option<"sparseEmitStrategy", "sparse-emit-strategy", "mlir::SparseEmitStrategy",
"mlir::SparseEmitStrategy::kFunctional",
"Emit functional code or interfaces (to debug) for sparse loops", [{llvm::cl::values(
>From 277b8de64d2e1a4bd9e3516b1ca2202af9d2a61c Mon Sep 17 00:00:00 2001
From: Yinying Li <yinyingli at google.com>
Date: Wed, 14 Aug 2024 23:57:17 +0000
Subject: [PATCH 2/2] Add parallelization strategy to sparsification and
bufferization pass.
---
.../Dialect/SparseTensor/Transforms/Passes.h | 3 +-
.../Pipelines/SparseTensorPipelines.cpp | 3 +-
.../SparsificationAndBufferizationPass.cpp | 13 ++++++--
.../SparseTensor/minipipeline_vector.mlir | 30 +++++++++++++++++++
4 files changed, 44 insertions(+), 5 deletions(-)
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/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
index c5eb965884396a..606cfa7094f32c 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_vector.mlir b/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
index 2475aa5139da48..0f5f3029e3b0ef 100755
--- a/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
+++ b/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
@@ -1,10 +1,14 @@
// RUN: mlir-opt %s --sparsification-and-bufferization | FileCheck %s --check-prefix=CHECK-NOVEC
// RUN: mlir-opt %s --sparsification-and-bufferization="vl=8" | FileCheck %s --check-prefix=CHECK-VEC
+// 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 optimization flags into
// the mini sparsification and bufferization pipeline.
#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+#SparseMatrix = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed, d1 : compressed)
+}>
#trait_sum_reduction = {
indexing_maps = [
@@ -15,6 +19,32 @@
doc = "x += SUM_i a(i)"
}
+#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-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>
+}
+
//
// CHECK-NOVEC-LABEL: func.func @sum_reduction
// CHECK-NOVEC: scf.for
More information about the Mlir-commits
mailing list