[Mlir-commits] [mlir] [mlir][affine] Fix crash in super-vectorizer when inner loop bounds depend on outer loop IV (PR #184618)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/184618

When the affine super-vectorizer processes a loop nest where an inner loop's bounds reference an outer loop's induction variable (IV), it incorrectly used the scalar outer loop's IV as operands when creating the vectorized inner loop. When the scalar loop nest was subsequently erased, those scalar IV values were destroyed while still being referenced by the vector inner loop's bounds, triggering a use_empty() assertion failure.

Fix this by using getScalarValueReplacementsFor() to map the scalar bound operands to their scalar replacements in the vector loop nest before creating the vectorized AffineForOp. This ensures the vector inner loop references the new vector loop's IV rather than the (soon-to-be-erased) scalar loop's IV.

Fixes #131135

>From 7ce7379a242f172f57d63bd9cb2b603ecfeb61e5 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 05:26:17 -0800
Subject: [PATCH] [mlir][affine] Fix crash in super-vectorizer when inner loop
 bounds depend on outer loop IV

When the affine super-vectorizer processes a loop nest where an inner loop's
bounds reference an outer loop's induction variable (IV), it incorrectly used
the scalar outer loop's IV as operands when creating the vectorized inner loop.
When the scalar loop nest was subsequently erased, those scalar IV values were
destroyed while still being referenced by the vector inner loop's bounds,
triggering a use_empty() assertion failure.

Fix this by using getScalarValueReplacementsFor() to map the scalar bound
operands to their scalar replacements in the vector loop nest before creating
the vectorized AffineForOp. This ensures the vector inner loop references the
new vector loop's IV rather than the (soon-to-be-erased) scalar loop's IV.

Fixes #131135
---
 .../Affine/Transforms/SuperVectorize.cpp      | 13 +++++++++---
 .../vectorize_inner_loop_dep_bounds.mlir      | 21 +++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)
 create mode 100644 mlir/test/Dialect/Affine/SuperVectorize/vectorize_inner_loop_dep_bounds.mlir

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 4e9f10b2c525c..f0728963b9636 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1383,10 +1383,17 @@ static Operation *vectorizeAffineForOp(AffineForOp forOp,
     }
   }
 
+  // Replace bound operands with scalar replacements from the vector loop nest.
+  // This is required when the bounds reference an outer loop's induction
+  // variable, which will be erased once the scalar loop nest is removed.
+  SmallVector<Value, 8> lbOperands, ubOperands;
+  state.getScalarValueReplacementsFor(forOp.getLowerBoundOperands(),
+                                      lbOperands);
+  state.getScalarValueReplacementsFor(forOp.getUpperBoundOperands(),
+                                      ubOperands);
   auto vecForOp = AffineForOp::create(
-      state.builder, forOp.getLoc(), forOp.getLowerBoundOperands(),
-      forOp.getLowerBoundMap(), forOp.getUpperBoundOperands(),
-      forOp.getUpperBoundMap(), newStep, vecIterOperands,
+      state.builder, forOp.getLoc(), lbOperands, forOp.getLowerBoundMap(),
+      ubOperands, forOp.getUpperBoundMap(), newStep, vecIterOperands,
       /*bodyBuilder=*/[](OpBuilder &, Location, Value, ValueRange) {
         // Make sure we don't create a default terminator in the loop body as
         // the proper terminator will be added during vectorization.
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_inner_loop_dep_bounds.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_inner_loop_dep_bounds.mlir
new file mode 100644
index 0000000000000..f6a201df9f512
--- /dev/null
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_inner_loop_dep_bounds.mlir
@@ -0,0 +1,21 @@
+// RUN: not mlir-opt %s -affine-super-vectorizer-test=vectorize-affine-loop-nest 2>&1 | FileCheck %s
+
+// Regression test for https://github.com/llvm/llvm-project/issues/131135
+// The vectorizer used to crash (assertion failure) when inner loop bounds
+// reference an outer loop induction variable. Verify a clean error is emitted.
+
+// CHECK: error:
+// CHECK-NOT: Assertion
+
+#map = affine_map<(d0) -> (d0)>
+#map1 = affine_map<(d0) -> (d0 + 1)>
+
+func.func @inner_loop_bounds_from_outer_iv(%arg0: memref<4x4xf32>, %arg1: memref<4xf32>) {
+  affine.for %i = 0 to 4 {
+    affine.for %j = #map(%i) to #map1(%i) {
+      %0 = affine.load %arg0[%j, %j] : memref<4x4xf32>
+      affine.store %0, %arg1[%j] : memref<4xf32>
+    }
+  }
+  return
+}