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

[wrengr]

Summary

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

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

52–78

do we have use cases for these options?

if not, I would be okay keeping this revision smaller and only add what you need?

mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py

122

not now, but probably in a later phase, let's move all those identical sparse compilers into a single place

[wrengr]

removing vector-to-scf options

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

52–78

I haven't seen any uses in the testsuite so far, but wasn't sure if that was an oversight or not

[wrengr]

rebase to resolve conflicts with D118648

[aartbik]

Are you waiting for us, or are you ready to commit this?

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

175

maybe we can call the second part "vectorization_options" and assign it before this, simply to make the functionality read easier.

also okay to do this in follow up revision where we share all the "SparseCompiler" classes between tests

[wrengr]

I was just waiting for the buildbot to finish, since I had to rebase it yesterday morning

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

 //===- Passes.h - Sparse tensor pipeline entry points -----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines prototypes of all sparse tensor pipelines.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 
 using namespace mlir::detail;
 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)};
   PassOptions::Option<int32_t> vectorization{
       *this, "vectorization-strategy", desc("Set the vectorization strategy"),
       init(0)};
   PassOptions::Option<int32_t> vectorLength{
       *this, "vl", desc("Set the vector length"), init(1)};
   PassOptions::Option<bool> enableSIMDIndex32{
       *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;
+  }
 };
 
 //===----------------------------------------------------------------------===//
 // Building and Registering.
 //===----------------------------------------------------------------------===//
 
 /// Adds the "sparse-compiler" pipeline to the `OpPassManager`.  This
 /// is the standard pipeline for taking sparsity-agnostic IR using

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // Pipeline implementation.
 //===----------------------------------------------------------------------===//
 
 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(createStdBufferizePass());
-  pm.addPass(mlir::bufferization::createFinalizingBufferizePass());
+  pm.addNestedPass<FuncOp>(createTensorBufferizePass());
+  pm.addNestedPass<FuncOp>(createStdBufferizePass());
+  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());
 }
 
 //===----------------------------------------------------------------------===//
 // Pipeline registration.
 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::registerSparseTensorPipelines() {
   PassPipelineRegistration<SparseCompilerOptions>(
       "sparse-compiler",
       "The standard pipeline for taking sparsity-agnostic IR using the"
       " sparse-tensor type, and lowering it to LLVM IR with concrete"
       " representations and algorithms for sparse tensors.",
       buildSparseCompiler);
 }

mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py

def build_compile_and_run_SDDMMM(attr: st.EncodingAttr, opt: str,
   expected[4, 1] = 3.0 * full_matmul[4, 1]
   c = rt.ranked_memref_to_numpy(mem_out[0])
   if np.allclose(c, expected):
     pass
   else:
     quit(f'FAILURE')
 
 
 class SparseCompiler:
   """Sparse compiler passes."""
 
   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,std-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):
     passmanager.PassManager.parse(self.pipeline).run(module)
 
 
 def main():
   support_lib = os.getenv('SUPPORT_LIB')

mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py

 
 
 class SparseCompiler:
   """Sparse compiler passes."""
 
   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,std-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):
     passmanager.PassManager.parse(self.pipeline).run(module)
 
 
 def main():
   support_lib = os.getenv('SUPPORT_LIB')

mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py

 """
 
 
 class _SparseCompiler:
   """Sparse compiler passes."""
 
   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,std-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)
 
 
 def _run_test(support_lib, kernel):
   """Compiles, runs and checks results."""
   module = ir.Module.parse(kernel)

mlir/test/Integration/Dialect/SparseTensor/python/test_output.py

 
 
 class SparseCompiler:
   """Sparse compiler passes."""
 
   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,std-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):
     passmanager.PassManager.parse(self.pipeline).run(module)
 
 
 def main():
   support_lib = os.getenv('SUPPORT_LIB')

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

 # other tests and programs.
 class SparseCompiler:
   """Sparse compiler passes."""
 
   def __init__(self, sparsification_options: str, support_lib: str):
     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,std-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)
 
   def __call__(self, module: ir.Module) -> ExecutionEngine:
     self._passmanager.run(module)
     return ExecutionEngine(module, opt_level=0, shared_libs=[self._support_lib])
 
 # ===----------------------------------------------------------------------=== #

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":BufferizationTransforms",
         ":ConversionPasses",
         ":LinalgTransforms",
         ":Pass",
         ":SparseTensor",
         ":SparseTensorTransforms",
         ":StandardOpsTransforms",
         ":TensorTransforms",
+        ":VectorToLLVM",
     ],
 )
 
 td_library(
     name = "StdOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/StandardOps/IR/Ops.td",
     ],