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

[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 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

---
Full diff: https://github.com/llvm/llvm-project/pull/184618.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 (+21) 


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

``````````

</details>


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