[MLIR] Fix crash in ValueBoundsConstraintSet for non-entry block args (llvm#185048)

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 llvm#119861

Assisted-by: Claude Code

Author: joker-eph

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)

mlir/test/Dialect/ControlFlow/value-bounds-op-interface-impl.mlir

@@ -0,0 +1,26 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds))' \
+// RUN:     -verify-diagnostics -split-input-file
+
+// Note: unstructured control flow (cf dialect) is not yet supported by the
+// ValueBoundsOpInterface. Block arguments from non-entry blocks cannot have
+// their bounds computed. The tests below verify that the infrastructure does
+// not crash on such inputs and fails gracefully instead.
+// See: https://github.com/llvm/llvm-project/issues/119861
+
+// Regression test: ValueBoundsConstraintSet must not crash when asked to
+// reify a bound for a non-entry block argument produced by unstructured
+// control flow.
+
+func.func @no_crash_non_entry_block_arg(%n: index) -> index {
+  %c0 = arith.constant 0 : index
+  cf.br ^bb1(%c0 : index)
+^bb1(%i: index):
+  // expected-error@+1 {{'test.reify_bound' op could not reify bound}}
+  %bound = "test.reify_bound"(%i) {type = "UB"} : (index) -> index
+  "test.some_use"(%bound) : (index) -> ()
+  %cond = arith.cmpi slt, %i, %n : index
+  %next = arith.addi %i, %c0 : index
+  cf.cond_br %cond, ^bb1(%next : index), ^bb2
+^bb2:
+  return %i : index
+}