diff --git a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h --- a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h +++ b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h @@ -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 @@ 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 { + // These options must be kept in sync with `SparsificationBase`. PassOptions::Option parallelization{ *this, "parallelization-strategy", desc("Set the parallelization strategy"), init(0)}; @@ -40,12 +42,53 @@ *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 reassociateFPReductions{ + *this, "reassociate-fp-reductions", + desc("Allows llvm to reassociate floating-point reductions for speed"), + init(false)}; + PassOptions::Option indexOptimizations{ + *this, "enable-index-optimizations", + desc("Allows compiler to assume indices fit in 32-bit if that yields " + "faster code"), + init(true)}; + PassOptions::Option amx{ + *this, "enable-amx", + desc("Enables the use of AMX dialect while lowering the vector dialect."), + init(false)}; + PassOptions::Option armNeon{*this, "enable-arm-neon", + desc("Enables the use of ArmNeon dialect " + "while lowering the vector dialect."), + init(false)}; + PassOptions::Option armSVE{*this, "enable-arm-sve", + desc("Enables the use of ArmSVE dialect " + "while lowering the vector dialect."), + init(false)}; + PassOptions::Option 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; + } }; //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp --- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp +++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp @@ -27,20 +27,22 @@ 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(createLinalgBufferizePass()); + pm.addNestedPass(createConvertLinalgToLoopsPass()); + pm.addNestedPass(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(createTensorBufferizePass()); + pm.addNestedPass( + mlir::bufferization::createFinalizingBufferizePass()); pm.addPass(createLowerAffinePass()); - pm.addPass(createConvertVectorToLLVMPass()); + pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions())); pm.addPass(createMemRefToLLVMPass()); - pm.addPass(createConvertMathToLLVMPass()); + pm.addNestedPass(createConvertMathToLLVMPass()); pm.addPass(createLowerToLLVMPass()); // --convert-std-to-llvm pm.addPass(createReconcileUnrealizedCastsPass()); } diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py +++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py @@ -124,18 +124,7 @@ 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): diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py +++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py @@ -114,18 +114,7 @@ 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): diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py +++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py @@ -66,18 +66,7 @@ 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) diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py +++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py @@ -74,18 +74,7 @@ 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): diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py +++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py @@ -172,18 +172,7 @@ 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) diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -1994,6 +1994,7 @@ ":SparseTensorTransforms", ":StandardOpsTransforms", ":TensorTransforms", + ":VectorToLLVM", ], )