[mlir][sparse] introduce sparse vectorization to the sparse compiler pipeline

[aartbik]

rebased with main (in particular, Bixia's buffer dep fix)

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

PassOptions::Option<mlir::SparseParallelizationStrategy> parallelization{
           clEnumValN(
               mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
               "any-storage-any-loop",
               "Enable sparse parallelization for any storage and loop."))};
 
   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{
       *this, "test-bufferization-analysis-only",
       desc("Run only the inplacability analysis"), init(false)};
 
   PassOptions::Option<bool> enableBufferInitialization{
       *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` differ 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"),
+  PassOptions::Option<bool> armSVE{
+      *this, "enable-arm-sve",
+      desc("Enables the use of ArmSVE dialect while lowering the vector "
+           "dialect"),
   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);
     opts.enableX86Vector(x86Vector);
     return opts;
   }
 };
 

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

 //===----------------------------------------------------------------------===//
 // Other rewriting rules and passes.
 //===----------------------------------------------------------------------===//
 
 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);
 
 std::unique_ptr<Pass> createSparseBufferRewritePass();
 std::unique_ptr<Pass>
 createSparseBufferRewritePass(bool enableBufferInitialization);
 

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

 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   pm.addNestedPass<func::FuncOp>(createLinalgGeneralizationPass());
   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());
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());
   // TODO(springerm): Add sparse support to the BufferDeallocation pass and add
   // it to this pipeline.
   pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());

mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp

if (auto load = cast.getDefiningOp<memref::LoadOp>()) {
           else if (etp.getIntOrFloatBitWidth() < 64 && !vl.enableSIMDIndex32)
             vload = rewriter.create<arith::ExtUIOp>(
                 loc, vectorType(vl, rewriter.getI64Type()), vload);
         }
         idxs.push_back(vload);
       }
       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;
 }
 
 #define UNAOP(xxx)                                                             \
   if (isa<xxx>(def)) {                                                         \
     if (codegen)                                                               \
if (vectorizeExpr(rewriter, forOp, vl, def->getOperand(0), codegen, vmask,
       TYPEDUNAOP(arith::FPToUIOp)
       TYPEDUNAOP(arith::SIToFPOp)
       TYPEDUNAOP(arith::UIToFPOp)
       TYPEDUNAOP(arith::ExtSIOp)
       TYPEDUNAOP(arith::ExtUIOp)
       TYPEDUNAOP(arith::IndexCastOp)
       TYPEDUNAOP(arith::TruncIOp)
       TYPEDUNAOP(arith::BitcastOp)
+      // TODO: complex?
     }
   } else if (def->getNumOperands() == 2) {
     Value vx, vy;
     if (vectorizeExpr(rewriter, forOp, vl, def->getOperand(0), codegen, vmask,
                       vx) &&
         vectorizeExpr(rewriter, forOp, vl, def->getOperand(1), codegen, vmask,
                       vy)) {
       BINOP(arith::MulFOp)
       BINOP(arith::MulIOp)
       BINOP(arith::DivFOp)
       BINOP(arith::DivSIOp)
       BINOP(arith::DivUIOp)
       BINOP(arith::AddFOp)
       BINOP(arith::AddIOp)
       BINOP(arith::SubFOp)
       BINOP(arith::SubIOp)
       BINOP(arith::AndIOp)
       BINOP(arith::OrIOp)
       BINOP(arith::XOrIOp)
+      // TODO: complex?
       // TODO: shift by invariant?
     }
   }
   return false;
 }
 
 #undef UNAOP
 #undef TYPEDUNAOP

 // Public method for populating vectorization rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with vectorization rules.
 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>,
                ReducChainRewriter<vector::BroadcastOp>>(patterns.getContext());
 }

mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #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;
 
 namespace mlir {
 namespace sparse_tensor {
 
 /// Return `true` if one of the given types is a sparse tensor type.

 class SparsificationAndBufferizationPass
     : public PassWrapper<SparsificationAndBufferizationPass,
                          OperationPass<ModuleOp>> {
 public:
   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
   /// `insertTensorCopies` in the form of `bufferization.alloc_tensor` ops.
   LogicalResult runDenseBufferization() {
     bufferization::OpFilter denseOpFilter;
     denseOpFilter.allowOperation([&](Operation *op) {
       if (containsSparseTensor(TypeRange(op->getResults())) ||
void runOnOperation() override {
 
     // Bufferize all sparse ops. No further analysis is needed. All required
     // buffer copies were already inserted by `insertTensorCopies` in the form
     // of `bufferization.alloc_tensor` ops.
     {
       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));
       } else {
         pm.addPass(createSparseTensorCodegenPass(enableBufferInitialization));
         pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
       }
       if (failed(runPipeline(pm, getOperation())))
         return signalPassFailure();
     }
 
     // Bufferize all dense ops.
     if (failed(runDenseBufferization()))
       signalPassFailure();
   }
 
 private:
   bufferization::OneShotBufferizationOptions bufferizationOptions;
   SparsificationOptions sparsificationOptions;
   SparseTensorConversionOptions sparseTensorConversionOptions;
   bool enableRuntimeLibrary;
   bool enableBufferInitialization;
+  unsigned vectorLength;
+  bool enableVLAVectorization;
+  bool enableSIMDIndex32;
 };
+
 } // namespace sparse_tensor
 } // namespace mlir
 
 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);
 }

mlir/test/Dialect/SparseTensor/sparse_vector_mv.mlir

+// 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>
+}