[mlir][sparse] expose optimization flags to mini pipeline #95158
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Some of the options only fed into the full sparse pipeline. However, some backends prefer to use the sparse minipipeline. This change exposes some important optimization flags to the pass as well. This prepares some SIMDization of PyTorch sparsified code.
aartbik
requested review from
PeimingLiu,
yinying-lisa-li and
matthias-springer
as code owners
|
@llvm/pr-subscribers-mlir Author: Aart Bik (aartbik) ChangesSome of the options only fed into the full sparse pipeline. However, some backends prefer to use the sparse minipipeline. This change exposes some important optimization flags to the pass as well. This prepares some SIMDization of PyTorch sparsified code. Full diff: https://github.com/llvm/llvm-project/pull/95158.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 2f844cee5ff52..7810ad2d60945 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -462,6 +462,18 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
"sparse_tensor::SparseTensorDialect",
"vector::VectorDialect"
];
+ // Important optimization options are made visible to the mini-pipeline
+ // so that clients can set these (when not using the full pipeline).
+ let options = [
+ Option<"vectorLength", "vl", "int32_t", "0",
+ "Set the vector length (use 0 to disable vectorization)">,
+ Option<"enableVLAVectorization", "enable-vla-vectorization", "bool", "false",
+ "Enable vector length agnostic vectorization">,
+ Option<"enableSIMDIndex32", "enable-simd-index32", "bool", "false",
+ "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">,
+ ];
}
#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index 3a8972072ac3b..13c750e83d045 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -61,20 +61,34 @@ class SparsificationAndBufferizationPass
: public impl::SparsificationAndBufferizationBase<
SparsificationAndBufferizationPass> {
public:
+ // Private pass options only.
SparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
bool createSparseDeallocs, bool enableRuntimeLibrary,
- bool enableBufferInitialization, unsigned vectorLength,
- bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen)
+ bool enableBufferInitialization)
: bufferizationOptions(bufferizationOptions),
sparsificationOptions(sparsificationOptions),
createSparseDeallocs(createSparseDeallocs),
enableRuntimeLibrary(enableRuntimeLibrary),
- enableBufferInitialization(enableBufferInitialization),
- vectorLength(vectorLength),
- enableVLAVectorization(enableVLAVectorization),
- enableSIMDIndex32(enableSIMDIndex32), enableGPULibgen(enableGPULibgen) {
+ enableBufferInitialization(enableBufferInitialization) {}
+ // Private pass options and visible pass options.
+ SparsificationAndBufferizationPass(
+ const bufferization::OneShotBufferizationOptions &bufferizationOptions,
+ const SparsificationOptions &sparsificationOptions,
+ bool createSparseDeallocs, bool enableRuntimeLibrary,
+ bool enableBufferInitialization, unsigned vl, bool vla, bool index32,
+ bool gpu)
+ : bufferizationOptions(bufferizationOptions),
+ sparsificationOptions(sparsificationOptions),
+ createSparseDeallocs(createSparseDeallocs),
+ enableRuntimeLibrary(enableRuntimeLibrary),
+ enableBufferInitialization(enableBufferInitialization) {
+ // Set the visible pass options explicitly.
+ vectorLength = vl;
+ enableVLAVectorization = vla;
+ enableSIMDIndex32 = index32;
+ enableGPULibgen = gpu;
}
/// Bufferize all dense ops. This assumes that no further analysis is needed
@@ -178,10 +192,6 @@ class SparsificationAndBufferizationPass
bool createSparseDeallocs;
bool enableRuntimeLibrary;
bool enableBufferInitialization;
- unsigned vectorLength;
- bool enableVLAVectorization;
- bool enableSIMDIndex32;
- bool enableGPULibgen;
};
} // namespace sparse_tensor
@@ -213,16 +223,13 @@ mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass() {
SparsificationOptions sparseOptions;
- return createSparsificationAndBufferizationPass(
+ return std::make_unique<
+ mlir::sparse_tensor::SparsificationAndBufferizationPass>(
getBufferizationOptionsForSparsification(/*analysisOnly=*/false),
sparseOptions,
/*createSparseDeallocs=*/false,
/*enableRuntimeLibrary=*/false,
- /*enableBufferInitialization=*/false,
- /*vectorLength=*/0,
- /*enableVLAVectorization=*/false,
- /*enableSIMDIndex32=*/false,
- /*enableGPULibgen=*/false);
+ /*enableBufferInitialization=*/false);
}
std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(
diff --git a/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir b/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
new file mode 100755
index 0000000000000..2475aa5139da4
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
@@ -0,0 +1,43 @@
+// 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
+
+// Test to ensure we can pass optimization flags into
+// the mini sparsification and bufferization pipeline.
+
+#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+#trait_sum_reduction = {
+ indexing_maps = [
+ affine_map<(i) -> (i)>, // a
+ affine_map<(i) -> ()> // x (scalar out)
+ ],
+ iterator_types = ["reduction"],
+ doc = "x += SUM_i a(i)"
+}
+
+//
+// CHECK-NOVEC-LABEL: func.func @sum_reduction
+// CHECK-NOVEC: scf.for
+// CHECK-NOVEC: arith.addf %{{.*}} %{{.*}} : f32
+// CHECK-NOVEC: }
+//
+// CHECK-VEC-LABEL: func.func @sum_reduction
+// CHECK-VEC: vector.insertelement
+// CHECK-VEC: scf.for
+// CHECK-VEC: vector.create_mask
+// CHECK-VEC: vector.maskedload
+// CHECK-VEC: arith.addf %{{.*}} %{{.*}} : vector<8xf32>
+// CHECK-VEC: }
+// CHECK-VEC: vector.reduction <add>
+//
+func.func @sum_reduction(%arga: tensor<?xf32, #SV>,
+ %argx: tensor<f32>) -> tensor<f32> {
+ %0 = linalg.generic #trait_sum_reduction
+ ins(%arga: tensor<?xf32, #SV>)
+ outs(%argx: tensor<f32>) {
+ ^bb(%a: f32, %x: f32):
+ %0 = arith.addf %x, %a : f32
+ linalg.yield %0 : f32
+ } -> tensor<f32>
+ return %0 : tensor<f32>
+}
|
|
@llvm/pr-subscribers-mlir-sparse Author: Aart Bik (aartbik) ChangesSome of the options only fed into the full sparse pipeline. However, some backends prefer to use the sparse minipipeline. This change exposes some important optimization flags to the pass as well. This prepares some SIMDization of PyTorch sparsified code. Full diff: https://github.com/llvm/llvm-project/pull/95158.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 2f844cee5ff52..7810ad2d60945 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -462,6 +462,18 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
"sparse_tensor::SparseTensorDialect",
"vector::VectorDialect"
];
+ // Important optimization options are made visible to the mini-pipeline
+ // so that clients can set these (when not using the full pipeline).
+ let options = [
+ Option<"vectorLength", "vl", "int32_t", "0",
+ "Set the vector length (use 0 to disable vectorization)">,
+ Option<"enableVLAVectorization", "enable-vla-vectorization", "bool", "false",
+ "Enable vector length agnostic vectorization">,
+ Option<"enableSIMDIndex32", "enable-simd-index32", "bool", "false",
+ "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">,
+ ];
}
#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index 3a8972072ac3b..13c750e83d045 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -61,20 +61,34 @@ class SparsificationAndBufferizationPass
: public impl::SparsificationAndBufferizationBase<
SparsificationAndBufferizationPass> {
public:
+ // Private pass options only.
SparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
bool createSparseDeallocs, bool enableRuntimeLibrary,
- bool enableBufferInitialization, unsigned vectorLength,
- bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen)
+ bool enableBufferInitialization)
: bufferizationOptions(bufferizationOptions),
sparsificationOptions(sparsificationOptions),
createSparseDeallocs(createSparseDeallocs),
enableRuntimeLibrary(enableRuntimeLibrary),
- enableBufferInitialization(enableBufferInitialization),
- vectorLength(vectorLength),
- enableVLAVectorization(enableVLAVectorization),
- enableSIMDIndex32(enableSIMDIndex32), enableGPULibgen(enableGPULibgen) {
+ enableBufferInitialization(enableBufferInitialization) {}
+ // Private pass options and visible pass options.
+ SparsificationAndBufferizationPass(
+ const bufferization::OneShotBufferizationOptions &bufferizationOptions,
+ const SparsificationOptions &sparsificationOptions,
+ bool createSparseDeallocs, bool enableRuntimeLibrary,
+ bool enableBufferInitialization, unsigned vl, bool vla, bool index32,
+ bool gpu)
+ : bufferizationOptions(bufferizationOptions),
+ sparsificationOptions(sparsificationOptions),
+ createSparseDeallocs(createSparseDeallocs),
+ enableRuntimeLibrary(enableRuntimeLibrary),
+ enableBufferInitialization(enableBufferInitialization) {
+ // Set the visible pass options explicitly.
+ vectorLength = vl;
+ enableVLAVectorization = vla;
+ enableSIMDIndex32 = index32;
+ enableGPULibgen = gpu;
}
/// Bufferize all dense ops. This assumes that no further analysis is needed
@@ -178,10 +192,6 @@ class SparsificationAndBufferizationPass
bool createSparseDeallocs;
bool enableRuntimeLibrary;
bool enableBufferInitialization;
- unsigned vectorLength;
- bool enableVLAVectorization;
- bool enableSIMDIndex32;
- bool enableGPULibgen;
};
} // namespace sparse_tensor
@@ -213,16 +223,13 @@ mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass() {
SparsificationOptions sparseOptions;
- return createSparsificationAndBufferizationPass(
+ return std::make_unique<
+ mlir::sparse_tensor::SparsificationAndBufferizationPass>(
getBufferizationOptionsForSparsification(/*analysisOnly=*/false),
sparseOptions,
/*createSparseDeallocs=*/false,
/*enableRuntimeLibrary=*/false,
- /*enableBufferInitialization=*/false,
- /*vectorLength=*/0,
- /*enableVLAVectorization=*/false,
- /*enableSIMDIndex32=*/false,
- /*enableGPULibgen=*/false);
+ /*enableBufferInitialization=*/false);
}
std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(
diff --git a/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir b/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
new file mode 100755
index 0000000000000..2475aa5139da4
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
@@ -0,0 +1,43 @@
+// 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
+
+// Test to ensure we can pass optimization flags into
+// the mini sparsification and bufferization pipeline.
+
+#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+#trait_sum_reduction = {
+ indexing_maps = [
+ affine_map<(i) -> (i)>, // a
+ affine_map<(i) -> ()> // x (scalar out)
+ ],
+ iterator_types = ["reduction"],
+ doc = "x += SUM_i a(i)"
+}
+
+//
+// CHECK-NOVEC-LABEL: func.func @sum_reduction
+// CHECK-NOVEC: scf.for
+// CHECK-NOVEC: arith.addf %{{.*}} %{{.*}} : f32
+// CHECK-NOVEC: }
+//
+// CHECK-VEC-LABEL: func.func @sum_reduction
+// CHECK-VEC: vector.insertelement
+// CHECK-VEC: scf.for
+// CHECK-VEC: vector.create_mask
+// CHECK-VEC: vector.maskedload
+// CHECK-VEC: arith.addf %{{.*}} %{{.*}} : vector<8xf32>
+// CHECK-VEC: }
+// CHECK-VEC: vector.reduction <add>
+//
+func.func @sum_reduction(%arga: tensor<?xf32, #SV>,
+ %argx: tensor<f32>) -> tensor<f32> {
+ %0 = linalg.generic #trait_sum_reduction
+ ins(%arga: tensor<?xf32, #SV>)
+ outs(%argx: tensor<f32>) {
+ ^bb(%a: f32, %x: f32):
+ %0 = arith.addf %x, %a : f32
+ linalg.yield %0 : f32
+ } -> tensor<f32>
+ return %0 : tensor<f32>
+}
|
PeimingLiu
approved these changes
Jun 11, 2024
yinying-lisa-li
approved these changes
Jun 11, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Some of the options only fed into the full sparse pipeline. However, some backends prefer to use the sparse minipipeline. This change exposes some important optimization flags to the pass as well. This prepares some SIMDization of PyTorch sparsified code.