[mlir][SparseTensor] Split scf.for loop into masked/unmasked parts

[springerm]

Summary

Apply the "for loop peeling" pattern from SCF dialect transforms. This pattern splits scf.for loops into full and partial iterations. In the full iteration, all masked loads/stores are canonicalized to unmasked loads/stores.

Depends On D107222

[aartbik]

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

9–10

you are running all the tests in this file now for VEC-2-PEELED
shall we test the peeling in a separate test that does just that?

e.g. sparse_vector_peeled.mlir with just that one function, the one RUN, and the one CHECK?

[springerm]

address comments

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

9–10

sounds good!

[aartbik]

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

139

would it be worthwhile to check this part too?

#map = affine_map<(d0, d1)[s0] -> (16, d0 - d1)>

or, if you don't want to do this here repeatedly, at least in the new sparse_vector_peeled test (but then as a before/after peeling, in that case we will need the CHECK-prefix again, see there

mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir

3

no need for prefix anymore, just remove and use CHECK

[springerm]

address comments

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

139

Yes, it should really be checked.

It just required a few smaller changes (replace CHECK-LABEL with -split-input-file).

[aartbik]

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

22

here and below, why did you remove the LABEL check? that should still work well with the split input feature, right?

mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir

21

probably still clean to have a

CHECK-LABEL: func @mul_ds

at the start

[springerm]

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

22

The affine map that I check for (#[[map]]) appears outside of the function body. If I use CHECK-LABEL here, I cannot refer to variables that are defined before the CHECK-LABEL line.

Only CHECK-VEC2 and CHECK-VEC3 have a map in some of their test cases, but I removed the LABEL part everywhere for consistency. I think -split-input-file has the same purpose as CHECK-LABEL: to avoid mixing up regex variables from different test cases.

mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir

21

Same problem here. If the CHECK-LABEL is before the CHECK-DAG: #[[map0..., I can no longer match the map0 (because it appears before the function). If I replace the CHECK: func @mul_s with a CHECK-LABEL: func @mul_s, the map was matched but I can no longer refer to [[map0]] after CHECK-LABEL.

But we should add a -split-input-file when more tests are added to this file.

[aartbik]

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

22

Looking at test/Dialect//Affine/canonicalize.mlir, there a split file is used, a CHECK-DAG outside func and a CHECK-LABEL. So it seems to work if you -DAG the affine map check?

[aartbik]

I will stop the bikeshedding and LG this revision. But please still have a look if you can keep the -LABEL.
If not, go ahead and I will have a look at it later once it goes in to convince myself one way or the other... ;-)

[nicolasvasilache]

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

22

FWIW, the trick here is to use the $ character: see e.g. mlir/test/Dialect/Linalg/tile-indexed.mlir

// TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)>
// TILE-10n25-LABEL: func @indexed_vector
...
// TILE-10n25:     %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[I]], %[[J]])

[aartbik]

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

22

FWIW, the trick here is to use the $ character: see e.g. mlir/test/Dialect/Linalg/tile-indexed.mlir

Ah sweet! Yes, I had seen that test/Dialect//Affine/canonicalize.mlir worked the way I wanted, for example, but overlooked that the $ was used there too. Thanks for the detail here!

[springerm]

address comments

[springerm]

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

22

Nice, that fixed the issue.

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

func @kernel_matvec(%arga: tensor<?x?xf64, #SparseMatrix>,
           linalg.yield %1 : f64
       } -> tensor<?xf64>
       return %0 : tensor<?xf64>
     }
     ```
   }];
   let constructor = "mlir::createSparsificationPass()";
   let dependentDialects = [
+    "AffineDialect",
     "LLVM::LLVMDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
     "vector::VectorDialect",
   ];
   let options = [
     Option<"parallelization", "parallelization-strategy", "int32_t", "0",

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

 //===- Sparsification.cpp - Implementation of sparsification --------------===//
 //
 // 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 file implements converting sparse tensor types to actual sparse code.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/SparseTensor/Utils/Merger.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/TensorEncoding.h"
 #include "llvm/ADT/SmallBitVector.h"
if (matchPattern(lo, m_Constant(&loInt)) &&
     if (((hiInt.getInt() - loInt.getInt()) % stepInt.getInt()) == 0)
       return rewriter.create<vector::BroadcastOp>(
           loc, mtp, rewriter.create<ConstantIntOp>(loc, 1, 1));
   }
   // Otherwise, generate a vector mask that avoids overrunning the upperbound
   // during vector execution. Here we rely on subsequent loop optimizations to
   // avoid executing the mask in all iterations, for example, by splitting the
   // loop into an unconditional vector loop and a scalar cleanup loop.
-  Value end = rewriter.create<SubIOp>(loc, hi, iv);
+  auto minMap = AffineMap::get(
+      /*dimCount=*/2, /*symbolCount=*/1,
+      {rewriter.getAffineSymbolExpr(0),
+       rewriter.getAffineDimExpr(0) - rewriter.getAffineDimExpr(1)},
+      rewriter.getContext());
+  Value end =
+      rewriter.createOrFold<AffineMinOp>(loc, minMap, ValueRange{hi, iv, step});
   return rewriter.create<vector::CreateMaskOp>(loc, mtp, end);
 }
 
 /// Generates a vectorized load lhs = a[ind[lo:hi]] or lhs = a[lo:hi].
 static Value genVectorLoad(CodeGen &codegen, PatternRewriter &rewriter,
                            Value ptr, ArrayRef<Value> args) {
   Location loc = ptr.getLoc();
   VectorType vtp = vectorType(codegen, ptr);

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

-// RUN: mlir-opt %s -sparsification="vectorization-strategy=0 vl=16" | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=0 vl=16" -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC0
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=1 vl=16" | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=1 vl=16" -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC1
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16" | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16" -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC2
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16 enable-simd-index32=true" | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16 enable-simd-index32=true" -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC3
 
 #DenseVector = #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>
 
-#SparseVector = #sparse_tensor.encoding<{
-  dimLevelType = [ "compressed" ],
-  pointerBitWidth = 32,
-  indexBitWidth = 32
-}>
-
-#SparseMatrix = #sparse_tensor.encoding<{
-  dimLevelType = [ "dense", "compressed" ],
-  pointerBitWidth = 32,
-  indexBitWidth = 32
-}>
-
 #trait_scale_d = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],
   iterator_types = ["parallel"],
   doc = "x(i) = a(i) * b"
 }
 
 //
 // CHECK-VEC0-LABEL: func @scale_d
 // CHECK-VEC0-DAG:   %[[c0:.*]] = constant 0 : index
 // CHECK-VEC0-DAG:   %[[c1:.*]] = constant 1 : index
 // CHECK-VEC0-DAG:   %[[c1024:.*]] = constant 1024 : index
 // CHECK-VEC0:       scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[c1]] {
 // CHECK-VEC0:         %[[l:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xf32>
 // CHECK-VEC0:         %[[m:.*]] = mulf %[[l]], %{{.*}} : f32
 // CHECK-VEC0:         store %[[m]], %{{.*}}[%[[i]]] : memref<1024xf32>
 // CHECK-VEC0:       }
%0 = linalg.generic #trait_scale_d
     outs(%argx: tensor<1024xf32>) {
       ^bb(%a: f32, %x: f32):
         %0 = mulf %a, %b : f32
         linalg.yield %0 : f32
   } -> tensor<1024xf32>
   return %0 : tensor<1024xf32>
 }
 
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed" ],
+  pointerBitWidth = 32,
+  indexBitWidth = 32
+}>
+
 #trait_mul_s = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>,  // b
     affine_map<(i) -> (i)>   // x (out)
   ],
   iterator_types = ["parallel"],
   doc = "x(i) = a(i) * b(i)"

 // CHECK-VEC1:         %[[ci:.*]] = index_cast %[[zi]] : i64 to index
 // CHECK-VEC1:         %[[la:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xf32>
 // CHECK-VEC1:         %[[lb:.*]] = memref.load %{{.*}}[%[[ci]]] : memref<1024xf32>
 // CHECK-VEC1:         %[[m:.*]] = mulf %[[la]], %[[lb]] : f32
 // CHECK-VEC1:         store %[[m]], %{{.*}}[%[[ci]]] : memref<1024xf32>
 // CHECK-VEC1:       }
 // CHECK-VEC1:       return
 //
+// CHECK-VEC2:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (16, d0 - d1)
 // CHECK-VEC2-LABEL: func @mul_s
 // CHECK-VEC2-DAG:   %[[c0:.*]] = constant 0 : index
 // CHECK-VEC2-DAG:   %[[c1:.*]] = constant 1 : index
 // CHECK-VEC2-DAG:   %[[c16:.*]] = constant 16 : index
 // CHECK-VEC2:       %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32>
 // CHECK-VEC2:       %[[a:.*]] = zexti %[[p]] : i32 to i64
 // CHECK-VEC2:       %[[q:.*]] = index_cast %[[a]] : i64 to index
 // CHECK-VEC2:       %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32>
 // CHECK-VEC2:       %[[b:.*]] = zexti %[[r]] : i32 to i64
 // CHECK-VEC2:       %[[s:.*]] = index_cast %[[b]] : i64 to index
 // CHECK-VEC2:       scf.for %[[i:.*]] = %[[q]] to %[[s]] step %[[c16]] {
-// CHECK-VEC2:         %[[sub:.*]] = subi %[[s]], %[[i]] : index
+// CHECK-VEC2:         %[[sub:.*]] = affine.min #[[$map]](%[[s]], %[[i]])[%[[c16]]]
 // CHECK-VEC2:         %[[mask:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
 // CHECK-VEC2:         %[[li:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32>
 // CHECK-VEC2:         %[[zi:.*]] = zexti %[[li]] : vector<16xi32> to vector<16xi64>
 // CHECK-VEC2:         %[[la:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 // CHECK-VEC2:         %[[lb:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 // CHECK-VEC2:         %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32>
 // CHECK-VEC2:         vector.scatter %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %[[m]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
 // CHECK-VEC2:       }
 // CHECK-VEC2:       return
 //
+// CHECK-VEC3:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (16, d0 - d1)
 // CHECK-VEC3-LABEL: func @mul_s
 // CHECK-VEC3-DAG:   %[[c0:.*]] = constant 0 : index
 // CHECK-VEC3-DAG:   %[[c1:.*]] = constant 1 : index
 // CHECK-VEC3-DAG:   %[[c16:.*]] = constant 16 : index
 // CHECK-VEC3:       %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32>
 // CHECK-VEC3:       %[[a:.*]] = zexti %[[p]] : i32 to i64
 // CHECK-VEC3:       %[[q:.*]] = index_cast %[[a]] : i64 to index
 // CHECK-VEC3:       %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32>
 // CHECK-VEC3:       %[[b:.*]] = zexti %[[r]] : i32 to i64
 // CHECK-VEC3:       %[[s:.*]] = index_cast %[[b]] : i64 to index
 // CHECK-VEC3:       scf.for %[[i:.*]] = %[[q]] to %[[s]] step %[[c16]] {
-// CHECK-VEC3:         %[[sub:.*]] = subi %{{.*}}, %[[i]] : index
+// CHECK-VEC3:         %[[sub:.*]] = affine.min #[[$map]](%[[s]], %[[i]])[%[[c16]]]
 // CHECK-VEC3:         %[[mask:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
 // 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:.*]] = 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
 //
 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 = mulf %a, %b : f32
         linalg.yield %0 : f32
   } -> tensor<1024xf32>
   return %0 : tensor<1024xf32>
 }
 
+// -----
+
+#DenseVector = #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>
+
 #trait_reduction_d = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>,  // b
     affine_map<(i) -> ()>    // x (out)
   ],
   iterator_types = ["reduction"],
   doc = "x += a(i) * b(i)"
outs(%argx: tensor<f32>) {
       ^bb(%a: f32, %b: f32, %x: f32):
         %0 = mulf %a, %b : f32
         %1 = addf %x, %0 : f32
         linalg.yield %1 : f32
   } -> tensor<f32>
   return %0 : tensor<f32>
 }
 
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  pointerBitWidth = 32,
+  indexBitWidth = 32
+}>
+
 #trait_mul_ds = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>,  // A
     affine_map<(i,j) -> (i,j)>,  // B
     affine_map<(i,j) -> (i,j)>   // X (out)
   ],
   iterator_types = ["parallel", "parallel"],
   doc = "X(i,j) = A(i,j) * B(i,j)"

 // CHECK-VEC1:           %[[la:.*]] = memref.load %{{.*}}[%[[j]]] : memref<?xf32>
 // CHECK-VEC1:           %[[lb:.*]] = memref.load %{{.*}}[%[[i]], %[[cj]]] : memref<512x1024xf32>
 // CHECK-VEC1:           %[[m:.*]] = mulf %[[la]], %[[lb]] : f32
 // CHECK-VEC1:           store %[[m]], %{{.*}}[%[[i]], %[[cj]]] : memref<512x1024xf32>
 // CHECK-VEC1:         }
 // CHECK-VEC1:       }
 // CHECK-VEC1:       return
 //
+// CHECK-VEC2:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (16, d0 - d1)
 // CHECK-VEC2-LABEL: func @mul_ds
 // CHECK-VEC2-DAG:   %[[c0:.*]] = constant 0 : index
 // CHECK-VEC2-DAG:   %[[c1:.*]] = constant 1 : index
 // CHECK-VEC2-DAG:   %[[c16:.*]] = constant 16 : index
 // CHECK-VEC2-DAG:   %[[c512:.*]] = constant 512 : index
 // CHECK-VEC2:       scf.for %[[i:.*]] = %[[c0]] to %[[c512]] step %[[c1]] {
 // CHECK-VEC2:         %[[p:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xi32>
 // CHECK-VEC2:         %[[a:.*]] = zexti %[[p]] : i32 to i64
 // CHECK-VEC2:         %[[q:.*]] = index_cast %[[a]] : i64 to index
 // CHECK-VEC2:         %[[a:.*]] = addi %[[i]], %[[c1]] : index
 // CHECK-VEC2:         %[[r:.*]] = memref.load %{{.*}}[%[[a]]] : memref<?xi32>
 // CHECK-VEC2:         %[[b:.*]] = zexti %[[r]] : i32 to i64
 // CHECK-VEC2:         %[[s:.*]] = index_cast %[[b]] : i64 to index
 // CHECK-VEC2:         scf.for %[[j:.*]] = %[[q]] to %[[s]] step %[[c16]] {
-// CHECK-VEC2:           %[[sub:.*]] = subi %[[s]], %[[j]] : index
+// CHECK-VEC2:           %[[sub:.*]] = affine.min #[[$map]](%[[s]], %[[j]])[%[[c16]]]
 // CHECK-VEC2:           %[[mask:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
 // CHECK-VEC2:           %[[lj:.*]] = vector.maskedload %{{.*}}[%[[j]]], %[[mask]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32>
 // CHECK-VEC2:           %[[zj:.*]] = zexti %[[lj]] : vector<16xi32> to vector<16xi64>
 // CHECK-VEC2:           %[[la:.*]] = vector.maskedload %{{.*}}[%[[j]]], %[[mask]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 // CHECK-VEC2:           %[[lb:.*]] = vector.gather %{{.*}}[%[[i]], %[[c0]]] [%[[zj]]], %[[mask]], %{{.*}} : memref<512x1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 // CHECK-VEC2:           %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32>
 // CHECK-VEC2:           vector.scatter %{{.*}}[%[[i]], %[[c0]]] [%[[zj]]], %[[mask]], %[[m]] : memref<512x1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
 // CHECK-VEC2:         }
 // CHECK-VEC2:       }
 // CHECK-VEC2:       return
 //
+// CHECK-VEC3:       #[[$map:.*]] = affine_map<(d0, d1)[s0] -> (16, d0 - d1)
 // CHECK-VEC3-LABEL: func @mul_ds
 // CHECK-VEC3-DAG:   %[[c0:.*]] = constant 0 : index
 // CHECK-VEC3-DAG:   %[[c1:.*]] = constant 1 : index
 // CHECK-VEC3-DAG:   %[[c16:.*]] = constant 16 : index
 // CHECK-VEC3-DAG:   %[[c512:.*]] = constant 512 : index
 // CHECK-VEC3:       scf.for %[[i:.*]] = %[[c0]] to %[[c512]] step %[[c1]] {
 // CHECK-VEC3:         %[[p:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xi32>
 // CHECK-VEC3:         %[[a:.*]] = zexti %[[p]] : i32 to i64
 // CHECK-VEC3:         %[[q:.*]] = index_cast %[[a]] : i64 to index
 // CHECK-VEC3:         %[[a:.*]] = addi %[[i]], %[[c1]] : index
 // CHECK-VEC3:         %[[r:.*]] = memref.load %{{.*}}[%[[a]]] : memref<?xi32>
 // CHECK-VEC3:         %[[b:.*]] = zexti %[[r]] : i32 to i64
 // CHECK-VEC3:         %[[s:.*]] = index_cast %[[b]] : i64 to index
 // CHECK-VEC3:         scf.for %[[j:.*]] = %[[q]] to %[[s]] step %[[c16]] {
-// CHECK-VEC3:           %[[sub:.*]] = subi %[[s]], %[[j]] : index
+// CHECK-VEC3:           %[[sub:.*]] = affine.min #[[$map]](%[[s]], %[[j]])[%[[c16]]]
 // CHECK-VEC3:           %[[mask:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
 // CHECK-VEC3:           %[[lj:.*]] = vector.maskedload %{{.*}}[%[[j]]], %[[mask]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32>
 // 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:.*]] = 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:       }

mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir

+// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16" -for-loop-peeling -canonicalize -split-input-file | \
+// RUN:   FileCheck %s
+
+#SparseVector = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed" ],
+  pointerBitWidth = 32,
+  indexBitWidth = 32
+}>
+
+#trait_mul_s = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> (i)>,  // b
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  iterator_types = ["parallel"],
+  doc = "x(i) = a(i) * b(i)"
+}
+
+// CHECK-DAG:   #[[$map0:.*]] = affine_map<()[s0, s1] -> (s0 + ((-s0 + s1) floordiv 16) * 16)>
+// CHECK-DAG:   #[[$map1:.*]] = affine_map<()[s0, s1] -> ((s0 - s1) mod 16)>
+// CHECK-LABEL: func @mul_s
+// CHECK-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-DAG:   %[[c1:.*]] = constant 1 : index
+// CHECK-DAG:   %[[c16:.*]] = constant 16 : index
+// CHECK:       %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32>
+// CHECK:       %[[a:.*]] = zexti %[[p]] : i32 to i64
+// CHECK:       %[[q:.*]] = index_cast %[[a]] : i64 to index
+// CHECK:       %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32>
+// CHECK:       %[[b:.*]] = zexti %[[r]] : i32 to i64
+// CHECK:       %[[s:.*]] = index_cast %[[b]] : i64 to index
+// CHECK:       %[[boundary:.*]] = affine.apply #[[$map0]]()[%[[q]], %[[s]]]
+// CHECK:       scf.for %[[i:.*]] = %[[q]] to %[[boundary]] step %[[c16]] {
+// CHECK:         %[[mask:.*]] = vector.constant_mask [16] : vector<16xi1>
+// CHECK:         %[[li:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xi32>, vector<16xi32>
+// CHECK:         %[[zi:.*]] = zexti %[[li]] : vector<16xi32> to vector<16xi64>
+// CHECK:         %[[la:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32>
+// CHECK:         %[[lb:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+// CHECK:         %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32>
+// CHECK:         vector.scatter %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %[[m]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
+// CHECK:       }
+// CHECK:       %[[has_more:.*]] = cmpi slt, %[[boundary]], %[[s]] : index
+// CHECK:       scf.if %[[has_more]] {
+// CHECK:         %[[sub:.*]] = affine.apply #[[$map1]]()[%[[s]], %[[q]]]
+// CHECK:         %[[mask2:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
+// CHECK:         %[[li2:.*]] = vector.maskedload %{{.*}}[%[[boundary]]], %[[mask2]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32>
+// CHECK:         %[[zi2:.*]] = zexti %[[li2]] : vector<16xi32> to vector<16xi64>
+// CHECK:         %[[la2:.*]] = vector.maskedload %{{.*}}[%[[boundary]]], %[[mask2]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+// CHECK:         %[[lb2:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+// CHECK:         %[[m2:.*]] = mulf %[[la2]], %[[lb2]] : vector<16xf32>
+// CHECK:         vector.scatter %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %[[m2]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
+// CHECK:       }
+// CHECK:       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 = mulf %a, %b : f32
+        linalg.yield %0 : f32
+  } -> tensor<1024xf32>
+  return %0 : tensor<1024xf32>
+}

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

 )
 
 cc_library(
     name = "SparseTensorTransforms",
     srcs = glob(["lib/Dialect/SparseTensor/Transforms/*.cpp"]),
     hdrs = ["include/mlir/Dialect/SparseTensor/Transforms/Passes.h"],
     includes = ["include"],
     deps = [
+        ":Affine",
         ":IR",
         ":LLVMDialect",
         ":LinalgOps",
         ":LinalgTransforms",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":SparseTensor",