[Mlir-commits] [mlir] 4161217 - [MLIR] fix: ValueBoundsConstraintSet invalidates queue (#218612)

[llvmlistbot at llvm.org]

Author: Alexey Kreshchuk
Date: 2026-08-25T07:54:24Z
New Revision: 4161217d470ade6232652d79d410083df9cc5c02

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

LOG: [MLIR] fix: ValueBoundsConstraintSet invalidates queue (#218612)

populateConstraints(AffineMap, ValueDimList) updates positionToValueDim,
changing the meaning of all elements in the worklist. In order to
preserve the correct mapping, I've changed the element type to ValueDim.
The new lit tests show the example where the issue happened.

Assisted-by: Cursor / Grok 4.6

Added: 
    

Modified: 
    mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
    mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
    mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
index f6dfa5855d011..50c2dfe64de2b 100644
--- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
+++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
@@ -460,7 +460,10 @@ class ValueBoundsConstraintSet
   DenseMap<ValueDim, int64_t> valueDimToPosition;
 
   /// Worklist of values/shape dimensions that have not been processed yet.
-  std::queue<int64_t> worklist;
+  /// Entries are identified by ValueDim, not by constraint-set column.
+  /// Inserting a SetDim shifts symbol columns, so a queued column index would
+  /// go stale.
+  std::queue<ValueDim> worklist;
 
   /// Constraint system of equalities and inequalities.
   FlatLinearConstraints cstr;

diff  --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
index 2ac5c8ffe46cb..d6d3b4d43cb55 100644
--- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
+++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
@@ -335,7 +335,7 @@ int64_t ValueBoundsConstraintSet::insert(Value value,
        cast<BlockArgument>(value).getOwner()->isEntryBlock())) {
     LDBG() << "Push to worklist: " << value
            << " (dim: " << dim.value_or(kIndexValue) << ")";
-    worklist.push(pos);
+    worklist.push(valueDim);
   }
 
   return pos;
@@ -411,11 +411,10 @@ bool ValueBoundsConstraintSet::isMapped(Value value,
 void ValueBoundsConstraintSet::processWorklist() {
   LDBG() << "Processing value bounds worklist...";
   while (!worklist.empty()) {
-    int64_t pos = worklist.front();
+    ValueDim valueDim = worklist.front();
     worklist.pop();
-    assert(positionToValueDim[pos].has_value() &&
-           "did not expect std::nullopt on worklist");
-    ValueDim valueDim = *positionToValueDim[pos];
+    assert(valueDimToPosition.contains(valueDim) &&
+           "expected mapped worklist entry");
     Value value = valueDim.first;
     int64_t dim = valueDim.second;
 

diff  --git a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
index beb80229bf857..7700e7b2144d6 100644
--- a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
@@ -327,3 +327,58 @@ func.func @affine_for_iv_multi_result_bounds(%n: index) {
   }
   return
 }
+
+// -----
+
+// Tiled `scf.for` + `affine.min` remainder. `%tile = minui(1, %n)` so
+// `%tile <= 1`, hence `%bytes = affine.min * 64` is `<= 64`. `{constant}`
+// reifies an exclusive UB, so the bound is 65.
+//
+// `affine.min` inserts `%iv`, `%ub`, and `%tile` as symbols. Processing
+// `%ub = arith.minui` calls `isProvablyNonNegative`, which inserts a new
+// SetDim and shifts every symbol column. The worklist must still visit
+// `%tile`; otherwise the bound is the remainder `16384 * 64` (exclusive
+// 1048577) and `tile <= 1` is never added.
+// CHECK-LABEL: func @affine_min_scf_for_stale_worklist_minui_tile
+//       CHECK:   %[[c65:.*]] = arith.constant 65 : index
+//       CHECK:   scf.yield %[[c65]]
+func.func @affine_min_scf_for_stale_worklist_minui_tile(%n: index) -> index {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c64 = arith.constant 64 : index
+  %c16384 = arith.constant 16384 : index
+  %ub = arith.minui %n, %c16384 : index
+  %tile = arith.minui %c1, %n : index
+  %ret = scf.for %iv = %c0 to %ub step %tile iter_args(%acc = %c0) -> index {
+    %min = affine.min affine_map<(d0)[s0, s1] -> (-d0 + s0, s1)>(%iv)[%ub, %tile]
+    %bytes = arith.muli %min, %c64 : index
+    %bound = "test.reify_bound"(%bytes) {type = "UB", constant} : (index) -> (index)
+    scf.yield %bound : index
+  }
+  return %ret : index
+}
+
+// -----
+
+// Same min, but the remainder is `arith.subi` + `arith.minui`. Tile is
+// processed, so the exclusive UB is 65 even before the worklist identity
+// fix.
+// CHECK-LABEL: func @affine_min_scf_for_minui_remainder
+//       CHECK:   %[[c65:.*]] = arith.constant 65 : index
+//       CHECK:   scf.yield %[[c65]]
+func.func @affine_min_scf_for_minui_remainder(%n: index) -> index {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c64 = arith.constant 64 : index
+  %c16384 = arith.constant 16384 : index
+  %ub = arith.minui %n, %c16384 : index
+  %tile = arith.minui %c1, %n : index
+  %ret = scf.for %iv = %c0 to %ub step %tile iter_args(%acc = %c0) -> index {
+    %rem = arith.subi %ub, %iv : index
+    %min = arith.minui %rem, %tile : index
+    %bytes = arith.muli %min, %c64 : index
+    %bound = "test.reify_bound"(%bytes) {type = "UB", constant} : (index) -> (index)
+    scf.yield %bound : index
+  }
+  return %ret : index
+}