[mlir][Sparse] Add option for VLA sparsification

[jsetoain]

Summary

Use "enable-vla-vectorization=vla" to generate a vector length agnostic
loops during vectorization. This option works for vectorization strategy 2.

[aartbik]

Thanks Javier for adding SVE to sparse!

Note that you may hold off on this revision until https://reviews.llvm.org/D117919 goes in (which restructures our pipeline a bit) to avoid conflicts later on....

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

98

period at end

[jsetoain]

In D118379#3276985, @aartbik wrote:

Note that you may hold off on this revision until https://reviews.llvm.org/D117919 goes in (which restructures our pipeline a bit) to avoid conflicts later on....

Yes, I've seen that one. I'm also waiting for solutions to the other two open patches, you can't lower to LLVM IR without those. And I also want to add more tests. I made it public while all those details are sorted out because I think this is a good reference point for VLA vectorization. It's an important milestone, and an interesting use case :-)

[jsetoain]

Added extra tests

[wrengr]

I just landed D117919. Please be sure to add the new enableVLAVectorization option to the SparseCompilerOptions in include/mlir/Dialect/SparseTensor/Pipelines/Passes.h as well.

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

52

I'd rather see this default specified on the no-parameter constructor below, rather than having it here; for the same reason we don't give defaults for any of the other parameters of this constructor.

[wrengr]

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

560

Why not rely on implicit coercion here?

[aartbik]

Note that we now have https://reviews.llvm.org/D118658 pending. Let's way until that goes in and then rebase and re-review

[wrengr]

Just landed D118658, if you want to rebase now

[aartbik]

@jsetoain can you plesae rebase and ping me for final review, thanks for your patience adjusting to the new pipeline!

[jsetoain]

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

560

General dislike for implicit conversions and a weak attempt at not obfuscating the semantic of "::get", but I see this is not discouraged in LLVM coding guidelines so I'm happy to change it for brevity's sake :-)

[jsetoain]

Update to latests changes in sparse compiler and address comments.

[jsetoain]

@aartbik Hi Aart, I believe this is ready for final review. Thank you!

[aartbik]

I love how simple you make it seem! A few nits, but after that good to go!

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

44 ↗

(On Diff #407481)

Just curious, is VLA the term going forward? Until now, SVE has been the norm (and the term scalable vector seems generic enough to include future platforms)?

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

99

No need to copy and store this boolean. We can always access it as codegen.options.enableVLAVectorization as we do for most other read-only settings. Most direct fields in codegen are read/write.

560

I am with you on this one, Javier. Can we find a compromise and use

unsigned numScalableDims = codegen.generateVLA;
return VectorType::get(codegen.curVecLength, etp, numScalableDims);

just to remind the reader that VectorType::get() expects a number, not a bool?

[jsetoain]

Address comments

[jsetoain]

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

44 ↗

(On Diff #407481)

VLA -> vector length agnostic, a vectorization model to handle scalable vectors.
SVE -> Scalable Vector Extension, Arm's specific implementation of scalable vectors.

Indeed, there are other, non-Arm, scalable vector ISAs (like RISCVV). I've tried to be consistent and call the programming model VLA (non-Arm-specific) and SVE (or Arm SVE) when it's specifically related to Arm's extension. A review of all related documentation is pending, so there might be inconsistencies here and there :-|

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

560

That is my preference as well. It makes it clear that it's taking the number of scalable dimensions, not whether it's scalable or not.

[aartbik]

\O/ yeah! sparse scalable vectorization! ship it!

[wrengr]

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

560

Fwiw, I too am not a fan of implicit conversions :) I only brought it up because the general mlir style seems to discourage making them explicit. I like the compromise

[jsetoain]

Adapt test to changes in vector.reduce

[jsetoain]

Fix missing case in type casting

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

enum class SparseVectorizationStrategy {
   kNone,
   kDenseInnerLoop,
   kAnyStorageInnerLoop
 };
 
 /// Sparsification options.
 struct SparsificationOptions {
   SparsificationOptions(SparseParallelizationStrategy p,
-                        SparseVectorizationStrategy v, unsigned vl, bool e)
+                        SparseVectorizationStrategy v, unsigned vl, bool e,
+                        bool vla = false)
       : parallelizationStrategy(p), vectorizationStrategy(v), vectorLength(vl),
-        enableSIMDIndex32(e) {}
+        enableSIMDIndex32(e), enableVLAVectorization(vla) {}
   SparsificationOptions()
       : SparsificationOptions(SparseParallelizationStrategy::kNone,
                               SparseVectorizationStrategy::kNone, 1u, false) {}
   SparseParallelizationStrategy parallelizationStrategy;
   SparseVectorizationStrategy vectorizationStrategy;
   unsigned vectorLength;
   bool enableSIMDIndex32;
+  bool enableVLAVectorization;
 };
 
 /// Sets up sparsification rewriting rules with the given options.
 void populateSparsificationPatterns(
     RewritePatternSet &patterns,
     const SparsificationOptions &options = SparsificationOptions());
 
 /// Sets up sparse tensor conversion rules.

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

def Sparsification : Pass<"sparsification", "ModuleOp"> {
   let options = [
     Option<"parallelization", "parallelization-strategy", "int32_t", "0",
            "Set the parallelization strategy">,
     Option<"vectorization", "vectorization-strategy", "int32_t", "0",
            "Set the vectorization strategy">,
     Option<"vectorLength", "vl", "int32_t", "1",
            "Set the vector length">,
     Option<"enableSIMDIndex32", "enable-simd-index32", "bool", "false",
-           "Enable i32 indexing into vectors (for efficiency)">
+           "Enable i32 indexing into vectors (for efficiency)">,
+    Option<"enableVLAVectorization", "enable-vla-vectorization", "bool",
+           "false", "Enable vector length agnostic vectorization">
   ];
 }
 
 def SparseTensorConversion : Pass<"sparse-tensor-conversion", "ModuleOp"> {
   let summary = "Apply conversion rules to sparse tensor primitives and types";
   let description = [{
     A pass that converts sparse tensor primitives to calls into a runtime
     support library. All sparse tensor types are converted into opaque

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

SparseVectorizationStrategy vectorOption() {
     }
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     // Translate strategy flags to strategy options.
     SparsificationOptions options(parallelOption(), vectorOption(),
-                                  vectorLength, enableSIMDIndex32);
+                                  vectorLength, enableSIMDIndex32,
+                                  enableVLAVectorization);
     // Apply rewriting.
     populateSparsificationPatterns(patterns, options);
     vector::populateVectorToVectorCanonicalizationPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
 class SparseTensorTypeConverter : public TypeConverter {

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

CodeGen(SparsificationOptions o, unsigned numTensors, unsigned numLoops,
       : options(o), loops(numLoops), sizes(numLoops), buffers(numTensors),
         pointers(numTensors, std::vector<Value>(numLoops)),
         indices(numTensors, std::vector<Value>(numLoops)),
         highs(numTensors, std::vector<Value>(numLoops)),
         pidxs(numTensors, std::vector<Value>(numLoops)),
         idxs(numTensors, std::vector<Value>(numLoops)), redExp(-1u), redVal(),
         redKind(kNoReduc), sparseOut(op), outerParNest(nest), lexIdx(),
         expValues(), expFilled(), expAdded(), expCount(), curVecLength(1),
-        curVecMask() {}
+        curVecMask(), generateVLA(false) {}
   /// Sparsification options.
   SparsificationOptions options;
   /// Universal dense indices and upper bounds (by index). The loops array
   /// is updated with the value of the universal dense index in the current
   /// loop. The sizes array is set once with the inferred dimension sizes.
   std::vector<Value> loops;
   std::vector<Value> sizes;
   /// Buffers for storing dense and sparse numerical values (by tensor).
struct CodeGen {
   Value lexIdx;
   Value expValues;
   Value expFilled;
   Value expAdded;
   Value expCount;
   // Current vector length and mask.
   unsigned curVecLength;
   Value curVecMask;
+  // Generate vector code in a vector length agnostic way
+  bool generateVLA;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler analysis methods.
 //===----------------------------------------------------------------------===//
 
if (!enc) {
       codegen.buffers[tensor] =
           rewriter.create<ToValuesOp>(loc, sparseTp, t->get());
     }
   }
 }
 
 /// Constructs vector type.
 static VectorType vectorType(CodeGen &codegen, Type etp) {
-  return VectorType::get(codegen.curVecLength, etp);
+  return VectorType::get(codegen.curVecLength, etp,
+                         codegen.generateVLA ? 1 : 0);
 }
 
 /// Constructs vector type from pointer.
 static VectorType vectorType(CodeGen &codegen, Value ptr) {
   return vectorType(codegen, ptr.getType().cast<MemRefType>().getElementType());
 }
 
 /// Constructs vector iteration mask.
static Operation *genFor(Merger &merger, CodeGen &codegen,
   bool isVector = !codegen.sparseOut &&
                   isVectorFor(codegen, isInner, isSparse) &&
                   denseUnitStrides(merger, op, idx);
   bool isParallel =
       !codegen.sparseOut &&
       isParallelFor(codegen, isOuter, isReduction, isSparse, isVector);
 
   // Prepare vector length.
-  if (isVector)
+  if (isVector) {
     codegen.curVecLength = codegen.options.vectorLength;
+    codegen.generateVLA = codegen.options.enableVLAVectorization;
+  }
 
   // Loop bounds and increment.
   Location loc = op.getLoc();
   Value lo = isSparse ? codegen.pidxs[tensor][idx] : codegen.loops[idx];
   Value hi = isSparse ? codegen.highs[tensor][idx] : codegen.sizes[idx];
   Value step = constantIndex(rewriter, loc, codegen.curVecLength);
+  if (codegen.generateVLA) {
+    Value vscale = rewriter.create<vector::VectorScaleOp>(
+        loc, IndexType::get(rewriter.getContext()));
+    step = rewriter.create<arith::MulIOp>(loc, vscale, step);
+  }
 
   // Emit a parallel loop.
   if (isParallel) {
     assert(!isVector);
     scf::ParallelOp parOp = rewriter.create<scf::ParallelOp>(loc, lo, hi, step);
     if (isSparse)
       codegen.pidxs[tensor][idx] = parOp.getInductionVars()[0];
     else

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

 // RUN: mlir-opt %s -sparsification="vectorization-strategy=0 vl=16" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC0
 // RUN: mlir-opt %s -sparsification="vectorization-strategy=1 vl=16" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC1
 // RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC2
 // RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16 enable-simd-index32=true" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC3
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=4 enable-vla-vectorization=true" -cse -split-input-file | \
+// RUN:   FileCheck %s --check-prefix=CHECK-VEC4
 
 #DenseVector = #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>
 
 #trait_scale_d = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],

 // CHECK-VEC2:       scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[c16]] {
 // CHECK-VEC2:         %[[r:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32>
 // CHECK-VEC2:         %[[b:.*]] = vector.broadcast %{{.*}} : f32 to vector<16xf32>
 // CHECK-VEC2:         %[[m:.*]] = arith.mulf %[[r]], %[[b]] : vector<16xf32>
 // CHECK-VEC2:         vector.store %[[m]], %{{.*}}[%[[i]]] : memref<1024xf32>, vector<16xf32>
 // CHECK-VEC2:       }
 // CHECK-VEC2:       return
 //
+// CHECK-VEC4:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)
+// CHECK-VEC4-LABEL: func @scale_d
+// CHECK-VEC4-DAG:   %[[c0:.*]] = arith.constant 0 : index
+// CHECK-VEC4-DAG:   %[[c4:.*]] = arith.constant 4 : index
+// CHECK-VEC4-DAG:   %[[c1024:.*]] = arith.constant 1024 : index
+// CHECK-VEC4-DAG:   %[[v0:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf32>
+// CHECK-VEC4-DAG:   %[[vscale:.*]] = vector.vscale
+// CHECK-VEC4:       %[[step:.*]] = arith.muli %[[vscale]], %[[c4]] : index
+// CHECK-VEC4:       scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[step]] {
+// CHECK-VEC4:         %[[sub:.*]] = affine.min #[[$map]](%[[c1024]], %[[i]])[%[[step]]]
+// CHECK-VEC4:         %[[mask:.*]] = vector.create_mask %[[sub]] : vector<[4]xi1>
+// CHECK-VEC4:         %[[val:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %[[v0]] : memref<?xf32>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+// CHECK-VEC4:         %[[scalev:.*]] = vector.broadcast %{{.*}} : f32 to vector<[4]xf32>
+// CHECK-VEC4:         %[[scaled:.*]] = arith.mulf %[[val]], %[[scalev]] : vector<[4]xf32>
+// CHECK-VEC4:         vector.maskedstore %{{.*}}[%[[i]]], %[[mask]], %[[scaled]] : memref<1024xf32>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK-VEC4:       }
+// CHECK-VEC4:       return
+//
 func @scale_d(%arga: tensor<1024xf32, #DenseVector>, %b: f32, %argx: tensor<1024xf32>) -> tensor<1024xf32> {
   %0 = linalg.generic #trait_scale_d
     ins(%arga: tensor<1024xf32, #DenseVector>)
     outs(%argx: tensor<1024xf32>) {
       ^bb(%a: f32, %x: f32):
         %0 = arith.mulf %a, %b : f32
         linalg.yield %0 : f32
   } -> tensor<1024xf32>

 // CHECK-VEC3:         %[[li:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32>
 // CHECK-VEC3:         %[[la:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 // CHECK-VEC3:         %[[lb:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[li]]], %[[mask]], %{{.*}} : memref<1024xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 // CHECK-VEC3:         %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : vector<16xf32>
 // CHECK-VEC3:         vector.scatter %{{.*}}[%[[c0]]] [%[[li]]], %[[mask]], %[[m]] : memref<1024xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>
 // CHECK-VEC3:       }
 // CHECK-VEC3:       return
 //
+// CHECK-VEC4:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)
+// CHECK-VEC4-LABEL: func @mul_s
+// CHECK-VEC4-DAG:   %[[c0:.*]] = arith.constant 0 : index
+// CHECK-VEC4-DAG:   %[[c1:.*]] = arith.constant 1 : index
+// CHECK-VEC4-DAG:   %[[c4:.*]] = arith.constant 4 : index
+// CHECK-VEC4-DAG:   %[[v0i:.*]] = arith.constant dense<0> : vector<[4]xi32>
+// CHECK-VEC4-DAG:   %[[v0f:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf32>
+// CHECK-VEC4:       %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32>
+// CHECK-VEC4:       %[[a:.*]] = arith.extui %[[p]] : i32 to i64
+// CHECK-VEC4:       %[[q:.*]] = arith.index_cast %[[a]] : i64 to index
+// CHECK-VEC4:       %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32>
+// CHECK-VEC4:       %[[b:.*]] = arith.extui %[[r]] : i32 to i64
+// CHECK-VEC4:       %[[s:.*]] = arith.index_cast %[[b]] : i64 to index
+// CHECK-VEC4:       %[[vscale:.*]] = vector.vscale
+// CHECK-VEC4:       %[[step:.*]] = arith.muli %[[vscale]], %[[c4]] : index
+// CHECK-VEC4:       scf.for %[[i:.*]] = %[[q]] to %[[s]] step %[[step]] {
+// CHECK-VEC4:         %[[sub:.*]] = affine.min #[[$map]](%[[s]], %[[i]])[%[[step]]]
+// CHECK-VEC4:         %[[mask:.*]] = vector.create_mask %[[sub]] : vector<[4]xi1>
+// CHECK-VEC4:         %[[li:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %[[v0i]] : memref<?xi32>, vector<[4]xi1>, vector<[4]xi32> into vector<[4]xi32>
+// CHECK-VEC4:         %[[lii64:.*]] = arith.extui %[[li]] : vector<[4]xi32> to vector<[4]xi64>
+// CHECK-VEC4:         %[[la:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %[[v0f]] : memref<?xf32>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+// CHECK-VEC4:         %[[lb:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[lii64]]], %[[mask]], %[[v0f]] : memref<1024xf32>, vector<[4]xi64>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+// CHECK-VEC4:         %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : vector<[4]xf32>
+// CHECK-VEC4:         vector.scatter %{{.*}}[%[[c0]]] [%[[lii64]]], %[[mask]], %[[m]] : memref<1024xf32>, vector<[4]xi64>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK-VEC4:       }
+// CHECK-VEC4:       return
+//
 func @mul_s(%arga: tensor<1024xf32, #SparseVector>, %argb: tensor<1024xf32>, %argx: tensor<1024xf32>) -> tensor<1024xf32> {
   %0 = linalg.generic #trait_mul_s
     ins(%arga, %argb: tensor<1024xf32, #SparseVector>, tensor<1024xf32>)
     outs(%argx: tensor<1024xf32>) {
       ^bb(%a: f32, %b: f32, %x: f32):
         %0 = arith.mulf %a, %b : f32
         linalg.yield %0 : f32
   } -> tensor<1024xf32>

 // CHECK-VEC2:         %[[lb:.*]] = vector.load %{{.*}}[%[[i]]] : memref<1024xf32>, vector<16xf32>
 // CHECK-VEC2:         %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : vector<16xf32>
 // CHECK-VEC2:         %[[a:.*]] = arith.addf %[[red_in]], %[[m]] : vector<16xf32>
 // CHECK-VEC2:         scf.yield %[[a]] : vector<16xf32>
 // CHECK-VEC2:       }
 // CHECK-VEC2:       %{{.*}} = vector.reduction "add", %[[red]] : vector<16xf32> into f32
 // CHECK-VEC2:       return
 //
+// CHECK-VEC4:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)
+// CHECK-VEC4-LABEL: func @reduction_d
+// CHECK-VEC4-DAG:   %[[c0:.*]] = arith.constant 0 : index
+// CHECK-VEC4-DAG:   %[[c4:.*]] = arith.constant 4 : index
+// CHECK-VEC4-DAG:   %[[c1024:.*]] = arith.constant 1024 : index
+// CHECK-VEC4-DAG:   %[[v0:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf32>
+// CHECK-VEC4:       %[[l:.*]] = memref.load %{{.*}}[] : memref<f32>
+// CHECK-VEC4:       %[[vscale:.*]] = vector.vscale
+// CHECK-VEC4:       %[[step:.*]] = arith.muli %[[vscale]], %[[c4]] : index
+// CHECK-VEC4:       %[[r:.*]] = vector.insertelement %[[l]], %[[v0]][%[[c0]] : index] : vector<[4]xf32>
+// CHECK-VEC4:       %[[red:.*]] = scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[step]] iter_args(%[[red_in:.*]] = %[[r]]) -> (vector<[4]xf32>) {
+// CHECK-VEC4:         %[[sub:.*]] = affine.min #[[$map]](%[[c1024]], %[[i]])[%[[step]]]
+// CHECK-VEC4:         %[[mask:.*]] = vector.create_mask %[[sub]] : vector<[4]xi1>
+// CHECK-VEC4:         %[[la:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %[[v0]] : memref<?xf32>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+// CHECK-VEC4:         %[[lb:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %[[v0]] : memref<1024xf32>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+// CHECK-VEC4:         %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : vector<[4]xf32>
+// CHECK-VEC4:         %[[a:.*]] = arith.addf %[[red_in]], %[[m]] : vector<[4]xf32>
+// CHECK-VEC4:         %[[sa:.*]] = select %[[mask]], %[[a]], %[[red_in]] : vector<[4]xi1>, vector<[4]xf32>
+// CHECK-VEC4:         scf.yield %[[sa]] : vector<[4]xf32>
+// CHECK-VEC4:       }
+// CHECK-VEC4:       %{{.*}} = vector.reduction "add", %[[red]] : vector<[4]xf32> into f32
+// CHECK-VEC4:       return
+//
 func @reduction_d(%arga: tensor<1024xf32, #DenseVector>, %argb: tensor<1024xf32>, %argx: tensor<f32>) -> tensor<f32> {
   %0 = linalg.generic #trait_reduction_d
     ins(%arga, %argb: tensor<1024xf32, #DenseVector>, tensor<1024xf32>)
     outs(%argx: tensor<f32>) {
       ^bb(%a: f32, %b: f32, %x: f32):
         %0 = arith.mulf %a, %b : f32
         %1 = arith.addf %x, %0 : f32
         linalg.yield %1 : f32

 // CHECK-VEC3:           %[[la:.*]] = vector.maskedload %{{.*}}[%[[j]]], %[[mask]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 // CHECK-VEC3:           %[[lb:.*]] = vector.gather %{{.*}}[%[[i]], %[[c0]]] [%[[lj]]], %[[mask]], %{{.*}} : memref<512x1024xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 // CHECK-VEC3:           %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : vector<16xf32>
 // CHECK-VEC3:           vector.scatter %{{.*}}[%[[i]], %[[c0]]] [%[[lj]]], %[[mask]], %[[m]] : memref<512x1024xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>
 // CHECK-VEC3:         }
 // CHECK-VEC3:       }
 // CHECK-VEC3:       return
 //
+// CHECK-VEC4:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)
+// CHECK-VEC4-LABEL: func @mul_ds
+// CHECK-VEC4-DAG:   %[[c0:.*]] = arith.constant 0 : index
+// CHECK-VEC4-DAG:   %[[c1:.*]] = arith.constant 1 : index
+// CHECK-VEC4-DAG:   %[[c4:.*]] = arith.constant 4 : index
+// CHECK-VEC4-DAG:   %[[c512:.*]] = arith.constant 512 : index
+// CHECK-VEC4-DAG:   %[[v0i:.*]] = arith.constant dense<0> : vector<[4]xi32>
+// CHECK-VEC4-DAG:   %[[v0f:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf32>
+// CHECK-VEC4:       scf.for %[[i:.*]] = %[[c0]] to %[[c512]] step %[[c1]] {
+// CHECK-VEC4:         %[[p:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xi32>
+// CHECK-VEC4:         %[[a:.*]] = arith.extui %[[p]] : i32 to i64
+// CHECK-VEC4:         %[[q:.*]] = arith.index_cast %[[a]] : i64 to index
+// CHECK-VEC4:         %[[a:.*]] = arith.addi %[[i]], %[[c1]] : index
+// CHECK-VEC4:         %[[r:.*]] = memref.load %{{.*}}[%[[a]]] : memref<?xi32>
+// CHECK-VEC4:         %[[b:.*]] = arith.extui %[[r]] : i32 to i64
+// CHECK-VEC4:         %[[s:.*]] = arith.index_cast %[[b]] : i64 to index
+// CHECK-VEC4:         %[[vscale:.*]] = vector.vscale
+// CHECK-VEC4:         %[[step:.*]] = arith.muli %[[vscale]], %[[c4]] : index
+// CHECK-VEC4:         scf.for %[[j:.*]] = %[[q]] to %[[s]] step %[[step]] {
+// CHECK-VEC4:           %[[sub:.*]] = affine.min #[[$map]](%[[s]], %[[j]])[%[[step]]]
+// CHECK-VEC4:           %[[mask:.*]] = vector.create_mask %[[sub]] : vector<[4]xi1>
+// CHECK-VEC4:           %[[lji32:.*]] = vector.maskedload %{{.*}}[%[[j]]], %[[mask]], %[[v0i]] : memref<?xi32>, vector<[4]xi1>, vector<[4]xi32> into vector<[4]xi32>
+// CHECK-VEC4:           %[[lj:.*]] = arith.extui %[[lji32]] : vector<[4]xi32> to vector<[4]xi64>
+// CHECK-VEC4:           %[[la:.*]] = vector.maskedload %{{.*}}[%[[j]]], %[[mask]], %[[v0f]] : memref<?xf32>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+// CHECK-VEC4:           %[[lb:.*]] = vector.gather %{{.*}}[%[[i]], %[[c0]]] [%[[lj]]], %[[mask]], %[[v0f]] : memref<512x1024xf32>, vector<[4]xi64>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+// CHECK-VEC4:           %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : vector<[4]xf32>
+// CHECK-VEC4:           vector.scatter %{{.*}}[%[[i]], %[[c0]]] [%[[lj]]], %[[mask]], %[[m]] : memref<512x1024xf32>, vector<[4]xi64>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK-VEC4:         }
+// CHECK-VEC4:       }
+// CHECK-VEC4:       return
+//
 func @mul_ds(%arga: tensor<512x1024xf32, #SparseMatrix>, %argb: tensor<512x1024xf32>, %argx: tensor<512x1024xf32>) -> tensor<512x1024xf32> {
   %0 = linalg.generic #trait_mul_ds
     ins(%arga, %argb: tensor<512x1024xf32, #SparseMatrix>, tensor<512x1024xf32>)
     outs(%argx: tensor<512x1024xf32>) {
       ^bb(%a: f32, %b: f32, %x: f32):
         %0 = arith.mulf %a, %b : f32
         linalg.yield %0 : f32
   } -> tensor<512x1024xf32>

 // CHECK-VEC2:           %[[x:.*]] = vector.gather %{{.*}}[%[[i1]], %[[c0]]] [%[[j]]], %[[mask]], %{{.*}} : memref<33x64xf64>
 // CHECK-VEC2:           %[[a:.*]] = vector.maskedload %{{.*}}[%[[jj]]], %[[mask]], %{{.*}} : memref<?xf64>
 // CHECK-VEC2:           %[[s:.*]] = arith.addf %[[x]], %[[a]] : vector<16xf64>
 // CHECK-VEC2:           vector.scatter %{{.*}}[%[[i1]], %[[c0]]] [%[[j]]], %[[mask]], %[[s]] : memref<33x64xf64>
 // CHECK-VEC2:         }
 // CHECK-VEC2:       }
 // CHECK-VEC2:       return
 //
+// CHECK-VEC4:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)
+// CHECK-VEC4-LABEL: func @add_dense
+// CHECK-VEC4-DAG:   %[[c0:.*]] = arith.constant 0 : index
+// CHECK-VEC4-DAG:   %[[c1:.*]] = arith.constant 1 : index
+// CHECK-VEC4-DAG:   %[[c4:.*]] = arith.constant 4 : index
+// CHECK-VEC4-DAG:   %[[c32:.*]] = arith.constant 32 : index
+// CHECK-VEC4-DAG:   %[[v0idx:.*]] = arith.constant dense<0> : vector<[4]xindex>
+// CHECK-VEC4-DAG:   %[[v0f64:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf64>
+// CHECK-VEC4:       scf.for %[[i:.*]] = %[[c0]] to %[[c32]] step %[[c1]] {
+// CHECK-VEC4:         %[[lo:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xindex>
+// CHECK-VEC4:         %[[i1:.*]] = arith.addi %[[i]], %[[c1]] : index
+// CHECK-VEC4:         %[[hi:.*]] = memref.load %{{.*}}[%[[i1]]] : memref<?xindex>
+// CHECK-VEC4:         %[[vscale:.*]] = vector.vscale
+// CHECK-VEC4:         %[[step:.*]] = arith.muli %[[vscale]], %[[c4]] : index
+// CHECK-VEC4:         scf.for %[[jj:.*]] = %[[lo]] to %[[hi]] step %[[step]] {
+// CHECK-VEC4:           %[[sub:.*]] = affine.min #[[$map]](%[[hi]], %[[jj]])[%[[step]]]
+// CHECK-VEC4:           %[[mask:.*]] = vector.create_mask %[[sub]] : vector<[4]xi1>
+// CHECK-VEC4:           %[[j:.*]] = vector.maskedload %{{.*}}[%[[jj]]], %[[mask]], %[[v0idx]] : memref<?xindex>
+// CHECK-VEC4:           %[[x:.*]] = vector.gather %{{.*}}[%[[i1]], %[[c0]]] [%[[j]]], %[[mask]], %[[v0f64]] : memref<33x64xf64>
+// CHECK-VEC4:           %[[a:.*]] = vector.maskedload %{{.*}}[%[[jj]]], %[[mask]], %[[v0f64]] : memref<?xf64>
+// CHECK-VEC4:           %[[s:.*]] = arith.addf %[[x]], %[[a]] : vector<[4]xf64>
+// CHECK-VEC4:           vector.scatter %{{.*}}[%[[i1]], %[[c0]]] [%[[j]]], %[[mask]], %[[s]] : memref<33x64xf64>
+// CHECK-VEC4:         }
+// CHECK-VEC4:       }
+// CHECK-VEC4:       return
+//
 func @add_dense(%arga: tensor<32x64xf64, #SparseMatrix>,
                 %argx: tensor<33x64xf64> {linalg.inplaceable = true}) -> tensor<33x64xf64> {
   %0 = linalg.generic #trait_affine
      ins(%arga: tensor<32x64xf64, #SparseMatrix>)
     outs(%argx: tensor<33x64xf64>) {
       ^bb(%a: f64, %x: f64):
         %0 = arith.addf %x, %a : f64
         linalg.yield %0 : f64
   } -> tensor<33x64xf64>
   return %0 : tensor<33x64xf64>
 }