replica
/
hanchenye-llvm-project
2
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[mlir][sparse] Updating sparse-compiler pipeline for python usage 

Explicitly nests passes for FuncOp, adds more options to the sparse-compiler pipeline, and updates python integration tests.  This should be sufficient to close https://github.com/llvm/llvm-project/issues/51751

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D118658
This commit is contained in:
wren romano 2022-02-01 14:19:11 -08:00
parent 35f7dd601d
commit 4998b1a6cd
8 changed files with 63 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
View File

@ -13,6 +13,7 @@
#ifndef MLIR_DIALECT_SPARSETENSOR_PIPELINES_PASSES_H_
#define MLIR_DIALECT_SPARSETENSOR_PIPELINES_PASSES_H_
#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Pass/PassOptions.h"
@ -22,12 +23,13 @@ using namespace llvm::cl;
namespace mlir {
namespace sparse_tensor {
/// Options for the "sparse-compiler" pipeline. So far this contains
/// only the same options as the sparsification pass, and must be kept
/// in sync with the `SparseTensor/Transforms/Passes.td` file. In the
/// future this may be extended with options for other passes in the pipeline.
/// Options for the "sparse-compiler" pipeline. So far this only contains
/// a subset of the options that can be set for the underlying passes,
/// because it must be manually kept in sync with the tablegen files
/// for those passes.
struct SparseCompilerOptions
: public PassPipelineOptions<SparseCompilerOptions> {
// These options must be kept in sync with `SparsificationBase`.
PassOptions::Option<int32_t> parallelization{
*this, "parallelization-strategy",
desc("Set the parallelization strategy"), init(0)};
@ -40,12 +42,53 @@ struct SparseCompilerOptions
*this, "enable-simd-index32",
desc("Enable i32 indexing into vectors (for efficiency)"), init(false)};
/// Projects out the options for the sparsification pass.
/// Projects out the options for `createSparsificationPass`.
SparsificationOptions sparsificationOptions() const {
return SparsificationOptions(sparseParallelizationStrategy(parallelization),
sparseVectorizationStrategy(vectorization),
vectorLength, enableSIMDIndex32);
}
// These options must be kept in sync with `ConvertVectorToLLVMBase`.
// TODO(wrengr): does `indexOptimizations` differ from `enableSIMDIndex32`?
PassOptions::Option<bool> reassociateFPReductions{
*this, "reassociate-fp-reductions",
desc("Allows llvm to reassociate floating-point reductions for speed"),
init(false)};
PassOptions::Option<bool> indexOptimizations{
*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."),
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."),
init(false)};
/// Projects out the options for `createConvertVectorToLLVMPass`.
LowerVectorToLLVMOptions lowerVectorToLLVMOptions() const {
LowerVectorToLLVMOptions opts{};
opts.enableReassociateFPReductions(reassociateFPReductions);
opts.enableIndexOptimizations(indexOptimizations);
opts.enableArmNeon(armNeon);
opts.enableArmSVE(armSVE);
opts.enableAMX(amx);
opts.enableX86Vector(x86Vector);
return opts;
}
};
//===----------------------------------------------------------------------===//

16
mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
View File

@ -27,20 +27,22 @@ using namespace mlir::sparse_tensor;
void mlir::sparse_tensor::buildSparseCompiler(
OpPassManager &pm, const SparseCompilerOptions &options) {
// TODO(wrengr): ensure the original `pm` is for ModuleOp
pm.addPass(createSparsificationPass(options.sparsificationOptions()));
pm.addPass(createSparseTensorConversionPass());
pm.addPass(createLinalgBufferizePass());
pm.addPass(createConvertLinalgToLoopsPass());
pm.addPass(createConvertVectorToSCFPass());
pm.addNestedPass<FuncOp>(createLinalgBufferizePass());
pm.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass());
pm.addNestedPass<FuncOp>(createConvertVectorToSCFPass());
pm.addPass(createLowerToCFGPass()); // --convert-scf-to-std
pm.addPass(createFuncBufferizePass());
pm.addPass(arith::createConstantBufferizePass());
pm.addPass(createTensorBufferizePass());
pm.addPass(mlir::bufferization::createFinalizingBufferizePass());
pm.addNestedPass<FuncOp>(createTensorBufferizePass());
pm.addNestedPass<FuncOp>(
mlir::bufferization::createFinalizingBufferizePass());
pm.addPass(createLowerAffinePass());
pm.addPass(createConvertVectorToLLVMPass());
pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
pm.addPass(createMemRefToLLVMPass());
pm.addPass(createConvertMathToLLVMPass());
pm.addNestedPass<FuncOp>(createConvertMathToLLVMPass());
pm.addPass(createLowerToLLVMPass()); // --convert-std-to-llvm
pm.addPass(createReconcileUnrealizedCastsPass());
}

13
mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py
View File

@ -124,18 +124,7 @@ class SparseCompiler:
def __init__(self, options: str):
pipeline = (
f'sparsification{{{options}}},'
f'sparse-tensor-conversion,'
f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
f'convert-scf-to-std,'
f'func-bufferize,'
f'arith-bufferize,'
f'builtin.func(tensor-bufferize,finalizing-bufferize),'
f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
f'lower-affine,'
f'convert-memref-to-llvm,'
f'convert-std-to-llvm,'
f'reconcile-unrealized-casts')
f'sparse-compiler{{{options} reassociate-fp-reductions=1 enable-index-optimizations=1}}')
self.pipeline = pipeline
def __call__(self, module: ir.Module):

13
mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
View File

@ -114,18 +114,7 @@ class SparseCompiler:
def __init__(self, options: str):
pipeline = (
f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
f'sparsification{{{options}}},'
f'sparse-tensor-conversion,'
f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
f'convert-scf-to-std,'
f'func-bufferize,'
f'arith-bufferize,'
f'builtin.func(tensor-bufferize,finalizing-bufferize),'
f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
f'lower-affine,'
f'convert-memref-to-llvm,'
f'convert-std-to-llvm,'
f'reconcile-unrealized-casts')
f'sparse-compiler{{{options} reassociate-fp-reductions=1 enable-index-optimizations=1}}')
self.pipeline = pipeline
def __call__(self, module: ir.Module):

13
mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py
View File

@ -66,18 +66,7 @@ class _SparseCompiler:
def __init__(self):
self.pipeline = (
f'sparsification,'
f'sparse-tensor-conversion,'
f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
f'convert-scf-to-std,'
f'func-bufferize,'
f'arith-bufferize,'
f'builtin.func(tensor-bufferize,finalizing-bufferize),'
f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
f'lower-affine,'
f'convert-memref-to-llvm,'
f'convert-std-to-llvm,'
f'reconcile-unrealized-casts')
f'sparse-compiler{{reassociate-fp-reductions=1 enable-index-optimizations=1}}')
def __call__(self, module: ir.Module):
passmanager.PassManager.parse(self.pipeline).run(module)

13
mlir/test/Integration/Dialect/SparseTensor/python/test_output.py
View File

@ -74,18 +74,7 @@ class SparseCompiler:
def __init__(self):
pipeline = (
f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
f'sparsification,'
f'sparse-tensor-conversion,'
f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
f'convert-scf-to-std,'
f'func-bufferize,'
f'arith-bufferize,'
f'builtin.func(tensor-bufferize,finalizing-bufferize),'
f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
f'lower-affine,'
f'convert-memref-to-llvm,'
f'convert-std-to-llvm,'
f'reconcile-unrealized-casts')
f'sparse-compiler{{reassociate-fp-reductions=1 enable-index-optimizations=1}}')
self.pipeline = pipeline
def __call__(self, module: ir.Module):

13
mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
View File

@ -172,18 +172,7 @@ class SparseCompiler:
self._support_lib = support_lib
self._pipeline = (
f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
f'sparsification{{{sparsification_options}}},'
f'sparse-tensor-conversion,'
f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
f'convert-scf-to-std,'
f'func-bufferize,'
f'arith-bufferize,'
f'builtin.func(tensor-bufferize,finalizing-bufferize),'
f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
f'lower-affine,'
f'convert-memref-to-llvm,'
f'convert-std-to-llvm,'
f'reconcile-unrealized-casts')
f'sparse-compiler{{{sparsification_options} reassociate-fp-reductions=1 enable-index-optimizations=1}}')
# Must be in the scope of a `with ir.Context():`
self._passmanager = PassManager.parse(self._pipeline)

1
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
View File

@ -1994,6 +1994,7 @@ cc_library(
":SparseTensorTransforms",
":StandardOpsTransforms",
":TensorTransforms",
":VectorToLLVM",
],
)