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

[Mehdi Amini]

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

When affine-super-vectorize creates the vectorized version of an inner loop whose bounds reference the outer loop's induction variable, it was passing the original scalar bound operands directly to the new vector AffineForOp. These operands reference block arguments of the scalar outer loop, which are destroyed when the scalar loop nest is erased, triggering a use_empty() assertion.

Fix this by replacing the bound operands with their scalar value replacements (via getScalarValueReplacementsFor) before creating the new vector loop, ensuring the vector inner loop uses the vector outer loop's induction variable.

Fixes #121321

>From 5e59b53195990f7978710a091a766efc3e00d02c Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 15:03:20 -0800
Subject: [PATCH] [mlir][affine] Fix crash in super-vectorize when inner loop
 bounds depend on outer IV

When affine-super-vectorize creates the vectorized version of an inner loop
whose bounds reference the outer loop's induction variable, it was passing
the original scalar bound operands directly to the new vector AffineForOp.
These operands reference block arguments of the scalar outer loop, which are
destroyed when the scalar loop nest is erased, triggering a use_empty()
assertion.

Fix this by replacing the bound operands with their scalar value replacements
(via getScalarValueReplacementsFor) before creating the new vector loop,
ensuring the vector inner loop uses the vector outer loop's induction variable.

Fixes #121321
---
 .../Affine/Transforms/SuperVectorize.cpp      | 13 +++++--
 .../vectorize_inner_loop_dep_bounds.mlir      | 38 +++++++++++++++++++
 2 files changed, 48 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..b4653cdd2a174 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 their scalar replacements. This is required
+  // when the bounds reference an outer loop's induction variable, which will
+  // be replaced (and eventually 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..4c64a481a940a
--- /dev/null
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_inner_loop_dep_bounds.mlir
@@ -0,0 +1,38 @@
+// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128" | FileCheck %s
+
+// Regression tests: vectorize inner loops whose bounds depend on the outer loop
+// induction variable. Previously this caused a crash due to a use-after-free
+// on the scalar outer loop's block argument when erasing the loop nest.
+
+#map = affine_map<(d0) -> (d0)>
+#map1 = affine_map<(d0) -> (d0 + 6)>
+
+// Both bounds depend on the outer IV.
+// CHECK-LABEL: func @inner_loop_dep_bounds
+// CHECK: affine.for %[[i:.*]] = 0 to 6 step 32
+// CHECK:   affine.for %[[j:.*]] = #map(%[[i]]) to #map1(%[[i]]) step 128
+// CHECK:     vector.transfer_write
+func.func @inner_loop_dep_bounds(%arg0: memref<32xf32>) {
+  %cst = arith.constant 1.0 : f32
+  affine.for %i = 0 to 6 step 32 {
+    affine.for %j = #map(%i) to #map1(%i) {
+      affine.store %cst, %arg0[%j] : memref<32xf32>
+    }
+  }
+  return
+}
+
+// Only upper bound depends on the outer IV; lower bound is a constant.
+// CHECK-LABEL: func @inner_loop_dep_upper_bound
+// CHECK: affine.for %[[i:.*]] = 0 to 64 step 32
+// CHECK:   affine.for %[[j:.*]] = 0 to #map1(%[[i]]) step 128
+// CHECK:     vector.transfer_write
+func.func @inner_loop_dep_upper_bound(%arg0: memref<64xf32>) {
+  %cst = arith.constant 1.0 : f32
+  affine.for %i = 0 to 64 step 32 {
+    affine.for %j = 0 to #map1(%i) {
+      affine.store %cst, %arg0[%j] : memref<64xf32>
+    }
+  }
+  return
+}