[Mlir-commits] [mlir] [MLIR] Fix crash in ValueBoundsConstraintSet for non-entry block args (PR #185048)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

When two vector transfer ops share a non-entry block argument as an index (e.g., in a loop with unstructured control flow), calling `ValueBoundsConstraintSet::areEqual` on those values caused a crash.

The first `populateConstraints` call would insert the block argument into the constraint set. The second call found it already mapped and called `getPos`, which hit an assert requiring the value to be either an OpResult or an entry-block argument.

Fix with two changes:
1. In `insert()`, suppress adding non-entry block arguments to the worklist. `ValueBoundsOpInterface` cannot derive bounds for such values, so the worklist push was a no-op and triggered the re-entrant `getPos` call.
2. Remove the overly conservative assert in `getPos`. Looking up a previously inserted non-entry block argument is valid; the assert was preventing legitimate use after the value had already been inserted.

Fixes #<!-- -->119861

Assisted-by: Claude Code

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


2 Files Affected:

- (modified) mlir/lib/Interfaces/ValueBoundsOpInterface.cpp (+10-4) 
- (modified) mlir/test/Dialect/Vector/vector-transferop-opt.mlir (+27) 


``````````diff
diff --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
index 85a156efd9474..35e3e92b91176 100644
--- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
+++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
@@ -282,7 +282,16 @@ int64_t ValueBoundsConstraintSet::insert(Value value,
     if (positionToValueDim[i].has_value())
       valueDimToPosition[*positionToValueDim[i]] = i;
 
-  if (addToWorklist) {
+  // Do not add block arguments from non-entry blocks to the worklist. The
+  // ValueBoundsOpInterface cannot derive any bounds for such values (they
+  // arise from unstructured control flow), so putting them on the worklist
+  // would be a no-op. More importantly, suppressing the worklist push ensures
+  // that processWorklist never calls getExpr on such a value a second time,
+  // which would otherwise cause the same value to be looked up as already
+  // mapped (triggering an unintended bug path).
+  if (addToWorklist &&
+      (!isa<BlockArgument>(value) ||
+       cast<BlockArgument>(value).getOwner()->isEntryBlock())) {
     LDBG() << "Push to worklist: " << value
            << " (dim: " << dim.value_or(kIndexValue) << ")";
     worklist.push(pos);
@@ -334,9 +343,6 @@ int64_t ValueBoundsConstraintSet::getPos(Value value,
                                          std::optional<int64_t> dim) const {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
-  assert((isa<OpResult>(value) ||
-          cast<BlockArgument>(value).getOwner()->isEntryBlock()) &&
-         "unstructured control flow is not supported");
 #endif // NDEBUG
   LDBG() << "Getting pos for: " << value
          << " (dim: " << dim.value_or(kIndexValue)
diff --git a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
index f4f7fb1ba0304..eb07df525a6e5 100644
--- a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
+++ b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
@@ -1,5 +1,7 @@
 // RUN: mlir-opt %s -test-vector-transferop-opt | FileCheck %s
 
+#leading_dim_1d_map = affine_map<(d0, d1) -> (d0)>
+
 // CHECK-LABEL: func @forward_dead_store
 //   CHECK-NOT:   vector.transfer_write
 //   CHECK-NOT:   vector.transfer_read
@@ -607,3 +609,28 @@ func.func @forward_and_eliminate_stores_through_trivial_aliases(
   vector.transfer_write %23, %subview_of_cast[%c0, %c0] {in_bounds = [true, true]} : vector<[8]x[8]xf32>, memref<?x?xf32, strided<[?, 1]>>
   return
 }
+
+// Regression test for https://github.com/llvm/llvm-project/issues/119861:
+// When two transfer ops in unstructured control flow share the same block
+// argument as an index, ValueBoundsConstraintSet must not crash when
+// isDisjointTransferIndices calls areEqual on those values.
+
+// CHECK-LABEL: func @no_crash_unstructured_cf_shared_index
+func.func @no_crash_unstructured_cf_shared_index(%alloc: memref<10x10xf32>) {
+  %v = arith.constant dense<0.0> : vector<4xf32>
+  %c0 = arith.constant 0 : index
+  cf.br ^bb1(%c0 : index)
+^bb1(%i: index):
+  %limit = arith.constant 10 : index
+  %cond = arith.cmpi slt, %i, %limit : index
+  cf.cond_br %cond, ^bb2, ^bb3
+^bb2:
+  // CHECK: vector.transfer_write
+  vector.transfer_write %v, %alloc[%i, %c0] {permutation_map = #leading_dim_1d_map} : vector<4xf32>, memref<10x10xf32>
+  // CHECK: vector.transfer_write
+  vector.transfer_write %v, %alloc[%i, %limit] {permutation_map = #leading_dim_1d_map} : vector<4xf32>, memref<10x10xf32>
+  %next = arith.addi %i, %limit : index
+  cf.br ^bb1(%next : index)
+^bb3:
+  return
+}

``````````

</details>


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