[Mlir-commits] [mlir] 16aa4e4 - [mlir][sparse] introduce sparse vectorization to the sparse compiler pipeline
Aart Bik
llvmlistbot at llvm.org
Wed Dec 7 16:07:03 PST 2022
Author: Aart Bik
Date: 2022-12-07T16:06:53-08:00
New Revision: 16aa4e4bd1754b6d7e964408c514fa0cce0678f5
URL: https://github.com/llvm/llvm-project/commit/16aa4e4bd1754b6d7e964408c514fa0cce0678f5
DIFF: https://github.com/llvm/llvm-project/commit/16aa4e4bd1754b6d7e964408c514fa0cce0678f5.diff
LOG: [mlir][sparse] introduce sparse vectorization to the sparse compiler pipeline
Reviewed By: Peiming
Differential Revision: https://reviews.llvm.org/D139581
Added:
mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir
Modified:
mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
index 5d51c5bd04403..4059cd4cc401a 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
@@ -56,7 +56,6 @@ struct SparseCompilerOptions
PassOptions::Option<bool> enableRuntimeLibrary{
*this, "enable-runtime-library",
desc("Enable runtime library for manipulating sparse tensors"),
- // TODO: Disable runtime library by default after feature complete.
init(true)};
PassOptions::Option<bool> testBufferizationAnalysisOnly{
@@ -67,56 +66,62 @@ struct SparseCompilerOptions
*this, "enable-buffer-initialization",
desc("Enable zero-initialization of memory buffers"), init(false)};
- /// Projects out the options for `createSparsificationPass`.
- SparsificationOptions sparsificationOptions() const {
- return SparsificationOptions(parallelization);
- }
+ PassOptions::Option<int32_t> vectorLength{
+ *this, "vl", desc("Set the vector length (0 disables vectorization)"),
+ init(0)};
// These options must be kept in sync with `SparseTensorConversionBase`.
PassOptions::Option<int32_t> sparseToSparse{
*this, "s2s-strategy",
desc("Set the strategy for sparse-to-sparse conversion"), init(0)};
- /// Projects out the options for `createSparsificationPass`.
- SparseTensorConversionOptions sparseTensorConversionOptions() const {
- return SparseTensorConversionOptions(
- sparseToSparseConversionStrategy(sparseToSparse));
- }
-
- // These options must be kept in sync with `ConvertVectorToLLVMBase`.
- // TODO(wrengr): does `indexOptimizations`
diff er from `enableSIMDIndex32`?
+ // These options must be kept in sync with the `ConvertVectorToLLVM`
+ // (defined in include/mlir/Dialect/SparseTensor/Pipelines/Passes.h).
PassOptions::Option<bool> reassociateFPReductions{
*this, "reassociate-fp-reductions",
desc("Allows llvm to reassociate floating-point reductions for speed"),
init(false)};
- PassOptions::Option<bool> indexOptimizations{
+ PassOptions::Option<bool> force32BitVectorIndices{
*this, "enable-index-optimizations",
desc("Allows compiler to assume indices fit in 32-bit if that yields "
"faster code"),
init(true)};
PassOptions::Option<bool> amx{
*this, "enable-amx",
- desc("Enables the use of AMX dialect while lowering the vector dialect."),
+ desc("Enables the use of AMX dialect while lowering the vector dialect"),
+ init(false)};
+ PassOptions::Option<bool> armNeon{
+ *this, "enable-arm-neon",
+ desc("Enables the use of ArmNeon dialect while lowering the vector "
+ "dialect"),
+ init(false)};
+ PassOptions::Option<bool> armSVE{
+ *this, "enable-arm-sve",
+ desc("Enables the use of ArmSVE dialect while lowering the vector "
+ "dialect"),
init(false)};
- PassOptions::Option<bool> armNeon{*this, "enable-arm-neon",
- desc("Enables the use of ArmNeon dialect "
- "while lowering the vector dialect."),
- init(false)};
- PassOptions::Option<bool> armSVE{*this, "enable-arm-sve",
- desc("Enables the use of ArmSVE dialect "
- "while lowering the vector dialect."),
- init(false)};
PassOptions::Option<bool> x86Vector{
*this, "enable-x86vector",
desc("Enables the use of X86Vector dialect while lowering the vector "
- "dialect."),
+ "dialect"),
init(false)};
+ /// Projects out the options for `createSparsificationPass`.
+ SparsificationOptions sparsificationOptions() const {
+ return SparsificationOptions(parallelization);
+ }
+
+ /// Projects out the options for `createSparseTensorConversionPass`.
+ SparseTensorConversionOptions sparseTensorConversionOptions() const {
+ return SparseTensorConversionOptions(
+ sparseToSparseConversionStrategy(sparseToSparse));
+ }
+
/// Projects out the options for `createConvertVectorToLLVMPass`.
LowerVectorToLLVMOptions lowerVectorToLLVMOptions() const {
LowerVectorToLLVMOptions opts{};
opts.enableReassociateFPReductions(reassociateFPReductions);
- opts.enableIndexOptimizations(indexOptimizations);
+ opts.enableIndexOptimizations(force32BitVectorIndices);
opts.enableArmNeon(armNeon);
opts.enableArmSVE(armSVE);
opts.enableAMX(amx);
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 252463566c414..d4e3210a6a342 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -166,7 +166,8 @@ std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions,
- bool enableRuntimeLibrary, bool enableBufferInitialization);
+ bool enableRuntimeLibrary, bool enableBufferInitialization,
+ unsigned vectorLength, bool enableVLAVectorization, bool enableSIMDIndex32);
void populateSparseBufferRewriting(RewritePatternSet &patterns,
bool enableBufferInitialization);
diff --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
index 326fed3d09de0..92fe090fb4e1f 100644
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -56,7 +56,10 @@ void mlir::sparse_tensor::buildSparseCompiler(
pm.addPass(createSparsificationAndBufferizationPass(
getBufferizationOptions(options.testBufferizationAnalysisOnly),
options.sparsificationOptions(), options.sparseTensorConversionOptions(),
- options.enableRuntimeLibrary, options.enableBufferInitialization));
+ options.enableRuntimeLibrary, options.enableBufferInitialization,
+ options.vectorLength,
+ /*enableVLAVectorization=*/options.armSVE,
+ /*enableSIMDIndex32=*/options.force32BitVectorIndices));
if (options.testBufferizationAnalysisOnly)
return;
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
index 55b4c9b2e485f..0668e82110037 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
@@ -282,6 +282,19 @@ static bool vectorizeSubscripts(PatternRewriter &rewriter, scf::ForOp forOp,
}
continue; // success so far
}
+ // Address calculation 'i = add inv, idx' (after LICM).
+ if (auto load = cast.getDefiningOp<arith::AddIOp>()) {
+ Value inv = load.getOperand(0);
+ Value idx = load.getOperand(1);
+ if (!inv.dyn_cast<BlockArgument>() &&
+ inv.getDefiningOp()->getBlock() != &forOp.getRegion().front() &&
+ idx.dyn_cast<BlockArgument>()) {
+ if (codegen)
+ idxs.push_back(
+ rewriter.create<arith::AddIOp>(forOp.getLoc(), inv, idx));
+ continue; // success so far
+ }
+ }
return false;
}
return true;
@@ -409,6 +422,7 @@ static bool vectorizeExpr(PatternRewriter &rewriter, scf::ForOp forOp, VL vl,
TYPEDUNAOP(arith::IndexCastOp)
TYPEDUNAOP(arith::TruncIOp)
TYPEDUNAOP(arith::BitcastOp)
+ // TODO: complex?
}
} else if (def->getNumOperands() == 2) {
Value vx, vy;
@@ -428,6 +442,7 @@ static bool vectorizeExpr(PatternRewriter &rewriter, scf::ForOp forOp, VL vl,
BINOP(arith::AndIOp)
BINOP(arith::OrIOp)
BINOP(arith::XOrIOp)
+ // TODO: complex?
// TODO: shift by invariant?
}
}
@@ -602,6 +617,7 @@ void mlir::populateSparseVectorizationPatterns(RewritePatternSet &patterns,
unsigned vectorLength,
bool enableVLAVectorization,
bool enableSIMDIndex32) {
+ assert(vectorLength > 0);
patterns.add<ForOpRewriter>(patterns.getContext(), vectorLength,
enableVLAVectorization, enableSIMDIndex32);
patterns.add<ReducChainRewriter<vector::InsertElementOp>,
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index 5633714ed807e..dd8dec9859568 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -18,6 +18,7 @@
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
using namespace mlir;
using namespace mlir::func;
@@ -53,12 +54,17 @@ class SparsificationAndBufferizationPass
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions,
- bool enableRuntimeLibrary, bool enableBufferInitialization)
+ bool enableRuntimeLibrary, bool enableBufferInitialization,
+ unsigned vectorLength, bool enableVLAVectorization,
+ bool enableSIMDIndex32)
: bufferizationOptions(bufferizationOptions),
sparsificationOptions(sparsificationOptions),
sparseTensorConversionOptions(sparseTensorConversionOptions),
enableRuntimeLibrary(enableRuntimeLibrary),
- enableBufferInitialization(enableBufferInitialization) {}
+ enableBufferInitialization(enableBufferInitialization),
+ vectorLength(vectorLength),
+ enableVLAVectorization(enableVLAVectorization),
+ enableSIMDIndex32(enableSIMDIndex32) {}
/// Bufferize all dense ops. This assumes that no further analysis is needed
/// and that all required buffer copies were already inserted by
@@ -127,6 +133,11 @@ class SparsificationAndBufferizationPass
OpPassManager pm("builtin.module");
pm.addPass(createSparsificationPass(sparsificationOptions));
pm.addPass(createPostSparsificationRewritePass(enableRuntimeLibrary));
+ if (vectorLength > 0) {
+ pm.addPass(mlir::createLoopInvariantCodeMotionPass());
+ pm.addPass(createSparseVectorizationPass(
+ vectorLength, enableVLAVectorization, enableSIMDIndex32));
+ }
if (enableRuntimeLibrary) {
pm.addPass(
createSparseTensorConversionPass(sparseTensorConversionOptions));
@@ -149,7 +160,11 @@ class SparsificationAndBufferizationPass
SparseTensorConversionOptions sparseTensorConversionOptions;
bool enableRuntimeLibrary;
bool enableBufferInitialization;
+ unsigned vectorLength;
+ bool enableVLAVectorization;
+ bool enableSIMDIndex32;
};
+
} // namespace sparse_tensor
} // namespace mlir
@@ -157,10 +172,13 @@ std::unique_ptr<Pass> mlir::createSparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions,
- bool enableRuntimeLibrary, bool enableBufferInitialization) {
+ bool enableRuntimeLibrary, bool enableBufferInitialization,
+ unsigned vectorLength, bool enableVLAVectorization,
+ bool enableSIMDIndex32) {
return std::make_unique<
mlir::sparse_tensor::SparsificationAndBufferizationPass>(
bufferizationOptions, sparsificationOptions,
sparseTensorConversionOptions, enableRuntimeLibrary,
- enableBufferInitialization);
+ enableBufferInitialization, vectorLength, enableVLAVectorization,
+ enableSIMDIndex32);
}
diff --git a/mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir
new file mode 100644
index 0000000000000..a43abe97b62e4
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s -sparse-compiler="vl=8" | FileCheck %s
+
+#Dense = #sparse_tensor.encoding<{
+ dimLevelType = [ "dense", "dense" ]
+}>
+
+#matvec = {
+ indexing_maps = [
+ affine_map<(i,j) -> (i,j)>, // A
+ affine_map<(i,j) -> (j)>, // b
+ affine_map<(i,j) -> (i)> // x (out)
+ ],
+ iterator_types = ["parallel", "reduction"],
+ doc = "X(i) += A(i,j) * B(j)"
+}
+
+// CHECK-LABEL: llvm.func @kernel_matvec
+// CHECK: llvm.intr.vector.reduce.fadd
+func.func @kernel_matvec(%arga: tensor<?x?xf32, #Dense>,
+ %argb: tensor<?xf32>,
+ %argx: tensor<?xf32>) -> tensor<?xf32> {
+ %x = linalg.generic #matvec
+ ins(%arga, %argb: tensor<?x?xf32, #Dense>, tensor<?xf32>)
+ outs(%argx: tensor<?xf32>) {
+ ^bb(%a: f32, %b: f32, %x: f32):
+ %0 = arith.mulf %a, %b : f32
+ %1 = arith.addf %x, %0 : f32
+ linalg.yield %1 : f32
+ } -> tensor<?xf32>
+ return %x : tensor<?xf32>
+}
More information about the Mlir-commits
mailing list