[Mlir-commits] [mlir] [MLIR][Affine] Fix null operands in simplifyConstrainedMinMaxOp (PR #189246)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/189246

`mlir::affine::simplifyConstrainedMinMaxOp` called `canonicalizeMapAndOperands` with `newOperands` that could contain null `Value()`s. These nulls came from `unpackOptionalValues(constraints.getMaybeValues(), newOperands)` where internal constraint variables added by `appendDimVar` (for `dimOp`, `dimOpBound`, and `resultDimStart*`) have no associated SSA values.

Passing null Values to `canonicalizeMapAndOperands` risks undefined behavior:
- `seenDims.find(null_value)` in the DenseMap causes all null operands to collide at the same key, producing incorrect dim remapping.
- Any null operand that remains referenced in the result map would propagate as a null Value into `AffineValueMap`, crashing callers that try to use those operands to create ops.

Fix: Before calling `canonicalizeMapAndOperands`, filter null operands from `newOperands` by replacing their dim/symbol positions in `newMap` with constant 0 (safe because internal constraint dims should not appear in the bound map expression) and compacting `newOperands` to contain only non-null Values.

Fixes #127436

Assisted-by: Claude Code

>From 790b9d95f1a3aadec6501e69a49ccf2d8a32a4ce Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Mar 2026 20:17:30 -0700
Subject: [PATCH] [MLIR][Affine] Fix null operands in
 simplifyConstrainedMinMaxOp

`mlir::affine::simplifyConstrainedMinMaxOp` called `canonicalizeMapAndOperands`
with `newOperands` that could contain null `Value()`s. These nulls came from
`unpackOptionalValues(constraints.getMaybeValues(), newOperands)` where internal
constraint variables added by `appendDimVar` (for `dimOp`, `dimOpBound`, and
`resultDimStart*`) have no associated SSA values.

Passing null Values to `canonicalizeMapAndOperands` risks undefined behavior:
- `seenDims.find(null_value)` in the DenseMap causes all null operands to
  collide at the same key, producing incorrect dim remapping.
- Any null operand that remains referenced in the result map would propagate
  as a null Value into `AffineValueMap`, crashing callers that try to use
  those operands to create ops.

Fix: Before calling `canonicalizeMapAndOperands`, filter null operands from
`newOperands` by replacing their dim/symbol positions in `newMap` with constant
0 (safe because internal constraint dims should not appear in the bound map
expression) and compacting `newOperands` to contain only non-null Values.

Fixes #127436

Assisted-by: Claude Code
---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp    | 37 +++++++++++++
 .../SCF/foreach-thread-canonicalization.mlir  | 53 ++++++++++++++++++-
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index ee893c14c9c6b..f5a02bb9309bd 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -2380,6 +2380,43 @@ FailureOr<AffineValueMap> mlir::affine::simplifyConstrainedMinMaxOp(
                               newMap.getNumDims(), newMap.getNumSymbols());
     }
   }
+
+  // Internal constraint variables (dimOp, dimOpBound, resultDimStart, etc.)
+  // have no associated SSA values (null Value()). Replace their corresponding
+  // dim/symbol positions in newMap with constant 0 and compact newOperands.
+  // These positions should be unreferenced in newMap (the bound was computed
+  // in terms of the original operands only), so replacing with 0 is safe.
+  // This prevents canonicalizeMapAndOperands from using null Values as
+  // DenseMap keys, which would cause undefined behavior.
+  {
+    unsigned numDims = newMap.getNumDims();
+    unsigned numSyms = newMap.getNumSymbols();
+    SmallVector<AffineExpr> dimReplacements(numDims), symReplacements(numSyms);
+    SmallVector<Value> filteredOperands;
+    filteredOperands.reserve(newOperands.size());
+    unsigned newDim = 0;
+    for (unsigned i = 0; i < numDims; ++i) {
+      if (newOperands[i]) {
+        dimReplacements[i] = getAffineDimExpr(newDim++, ctx);
+        filteredOperands.push_back(newOperands[i]);
+      } else {
+        dimReplacements[i] = getAffineConstantExpr(0, ctx);
+      }
+    }
+    unsigned newSym = 0;
+    for (unsigned i = 0; i < numSyms; ++i) {
+      if (newOperands[numDims + i]) {
+        symReplacements[i] = getAffineSymbolExpr(newSym++, ctx);
+        filteredOperands.push_back(newOperands[numDims + i]);
+      } else {
+        symReplacements[i] = getAffineConstantExpr(0, ctx);
+      }
+    }
+    newMap = newMap.replaceDimsAndSymbols(dimReplacements, symReplacements,
+                                          newDim, newSym);
+    newOperands = std::move(filteredOperands);
+  }
+
   affine::canonicalizeMapAndOperands(&newMap, &newOperands);
   return AffineValueMap(newMap, newOperands);
 }
diff --git a/mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir b/mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir
index 5b65c49ea6ed1..9d0c65e06d360 100644
--- a/mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir
+++ b/mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir
@@ -1,5 +1,6 @@
-// RUN: mlir-opt %s -scf-for-loop-canonicalization | FileCheck %s
+// RUN: mlir-opt %s -scf-for-loop-canonicalization -split-input-file | FileCheck %s
 
+// CHECK-LABEL: func @reduce
 func.func @reduce() {
   // CHECK: %[[C64:.*]] = arith.constant 64 : index
   %c2 = arith.constant 2 : index
@@ -34,3 +35,53 @@ func.func @reduce() {
   }
   return
 }
+
+// -----
+
+// Regression test for GH#127436: simplifyConstrainedMinMaxOp must handle
+// null-value operands produced when scf.forall has static (integer attribute)
+// bounds. In addLoopRangeConstraints, static lb/ub produce null-value symbol
+// variables in FlatAffineValueConstraints (via appendSymbolVar with no SSA
+// value). These must be removed before calling canonicalizeMapAndOperands to
+// avoid undefined behavior (null Value used as DenseMap key or passed to
+// matchPattern). The filter in simplifyConstrainedMinMaxOp compacts the
+// operands by replacing null positions with constant 0 and discarding them.
+
+// CHECK-LABEL: func @forall_static_bounds_affine_min_simplify
+// CHECK: %[[C64:.*]] = arith.constant 64 : i32
+// CHECK-NOT: affine.min
+// CHECK: memref.store %[[C64]]
+func.func @forall_static_bounds_affine_min_simplify(%A : memref<128xi32>) {
+  // lb=0, ub=2, step=1 are static integers -> IntegerAttr in OpFoldResult
+  // -> null-value symbols in the constraint system.
+  scf.forall (%i) = (0) to (2) step (1) {
+    // For i in [0, 2): min(-64*i + 128, 64) = 64 always.
+    %tile = affine.min affine_map<(d0) -> (d0 * -64 + 128, 64)>(%i)
+    %cast = arith.index_cast %tile : index to i32
+    %c0 = arith.constant 0 : index
+    memref.store %cast, %A[%c0] : memref<128xi32>
+    scf.forall.in_parallel {}
+  }
+  return
+}
+
+// -----
+
+// Regression test for GH#127436: same as above but with a 3-result affine.min
+// to exercise multiple resultDimStart null-value dim slots.
+
+// CHECK-LABEL: func @forall_static_bounds_multi_result_min_simplify
+// CHECK: %[[C32:.*]] = arith.constant 32 : i32
+// CHECK-NOT: affine.min
+// CHECK: memref.store %[[C32]]
+func.func @forall_static_bounds_multi_result_min_simplify(%A : memref<128xi32>) {
+  scf.forall (%i) = (0) to (2) step (1) {
+    // 3 results; for i in [0, 2): all results are >= 32, so min = 32.
+    %tile = affine.min affine_map<(d0) -> (d0 * -64 + 128, d0 * -64 + 96, 32)>(%i)
+    %cast = arith.index_cast %tile : index to i32
+    %c0 = arith.constant 0 : index
+    memref.store %cast, %A[%c0] : memref<128xi32>
+    scf.forall.in_parallel {}
+  }
+  return
+}