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

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-04-03T10:17:50+02:00
New Revision: d725513e7d92a2f815eb19c3aee06b599b42e834

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

LOG: [MLIR][Affine] Fix null operands in simplifyConstrainedMinMaxOp (#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

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Analysis/Utils.cpp
    mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index ee893c14c9c6b..ebe932a14694a 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -2380,6 +2380,47 @@ 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 {
+        assert(!newMap.isFunctionOfDim(i) &&
+               "null-valued dim operand referenced in bound map");
+        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 {
+        assert(!newMap.isFunctionOfSymbol(i) &&
+               "null-valued symbol operand referenced in bound map");
+        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
+}