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

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-06T12:03:30+01:00
New Revision: c044b8b59d8e27070f1ebe4d66733de498a86efc

URL: https://github.com/llvm/llvm-project/commit/c044b8b59d8e27070f1ebe4d66733de498a86efc
DIFF: https://github.com/llvm/llvm-project/commit/c044b8b59d8e27070f1ebe4d66733de498a86efc.diff

LOG: [mlir][affine] Fix crash in super-vectorize when inner loop bounds depend on outer IV (#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

Added: 
    mlir/test/Dialect/Affine/SuperVectorize/vectorize_inner_loop_dep_bounds.mlir

Modified: 
    mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

Removed: 
    


################################################################################
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
+}