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 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 dimReplacements(numDims), symReplacements(numSyms); + SmallVector 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 +}