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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/227670.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp (+8-6) 
- (modified) mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir (+16) 


``````````diff
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
+}

``````````

</details>


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