[Mlir-commits] [mlir] [mlir][affine] Do not super-vectorize loops with a non-unit step (PR #227670)

[Samarth Narang]

https://github.com/snarang181 created https://github.com/llvm/llvm-project/pull/227670

The vectorizer multiplies the step of the vectorized loop by the vectorization factor but vectorizes the body with contiguous transfers, so for a loop with a non-unit step the lanes cover iv, iv + 1, ... instead of iv, iv + step, ... and stores land on the wrong elements. The existing guard only rejected this for reduction loops; we now reject it for every loop along the vectorized dimension.

Fixes #225247

>From 97cf8aceb717015c04a1941315b06d9a4af99076 Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Wed, 30 Sep 2026 11:46:21 +0000
Subject: [PATCH] [mlir][affine] Do not super-vectorize loops with a non-unit
 step

The vectorizer multiplies the step of the vectorized loop by the vectorization
factor but vectorizes the body with contiguous transfers, so for a loop with a
non-unit step the lanes cover iv, iv + 1, ... instead of iv, iv + step, ...
and stores land on the wrong elements. The existing guard only rejected this
for reduction loops; reject it for every loop along the vectorized dimension.

Fixes #225247
---
 .../Dialect/Affine/Transforms/SuperVectorize.cpp | 14 ++++++++------
 .../Affine/SuperVectorize/vectorize_1d.mlir      | 16 ++++++++++++++++
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 6c6e8ed979d23..b3a98ead0e74f 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1360,12 +1360,14 @@ static Operation *vectorizeAffineForOp(AffineForOp forOp,
   auto loopToVecDimIt = strategy.loopToVectorDim.find(forOp);
   bool isLoopVecDim = loopToVecDimIt != strategy.loopToVectorDim.end();
 
-  // TODO: Vectorization of reduction loops is not supported for non-unit steps.
-  if (isLoopVecDim && forOp.getNumIterOperands() > 0 && forOp.getStep() != 1) {
-    LLVM_DEBUG(
-        dbgs()
-        << "\n[early-vect]+++++ unsupported step size for reduction loop: "
-        << forOp.getStep() << "\n");
+  // The body is vectorized with contiguous transfers while the vector loop
+  // step is the original step times the vectorization factor, which is only
+  // consistent for a unit step.
+  // TODO: Support non-unit steps with strided accesses.
+  if (isLoopVecDim && forOp.getStep() != 1) {
+    LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ unsupported step size for "
+                         "vectorized loop: "
+                      << forOp.getStep() << "\n");
     return nullptr;
   }
 
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
index f99190be6365d..167ab5f354422 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir
@@ -722,3 +722,19 @@ func.func @index_const_inside_loop(%mem: memref<1x8xi32>) {
   }
   return
 }
+
+// -----
+
+// A non-unit loop step cannot be vectorized as contiguous lanes: the scalar
+// loop touches iv, iv + step, ... while a plain transfer touches iv, iv + 1, ...
+// CHECK-LABEL: @non_unit_step_not_vectorized
+// CHECK:       affine.for %{{.*}} = -4 to 13 step 3
+// CHECK:         affine.store
+// CHECK-NOT:   vector.transfer_write
+func.func @non_unit_step_not_vectorized(%out: memref<32xi32>) {
+  %one = arith.constant 1 : i32
+  affine.for %iv = -4 to 13 step 3 {
+    affine.store %one, %out[%iv + 8] : memref<32xi32>
+  }
+  return
+}