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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp (+10-3) 
- (added) mlir/test/Dialect/Affine/SuperVectorize/vectorize_inner_loop_dep_bounds.mlir (+38) 


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

``````````

</details>


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