[Mlir-commits] [mlir] [MLIR][SCFToAffine] Validate inlined affine bound operands (PR #217473)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/217473

>From 7f2ec6e70723ac8fd92618bc3df1d46232289a0e Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 19 Aug 2026 07:42:43 -0700
Subject: [PATCH] [MLIR][SCFToAffine] Validate inlined affine bound operands

Validate affine.max and affine.min operands before inlining their maps into
affine.for bounds. This defers nested loops until enclosing induction
variables become valid affine dimensions and keeps the exported pattern
independent of greedy traversal order.

Assisted-by: Codex
---
 .../Conversion/SCFToAffine/SCFToAffine.cpp    | 32 +++++++++++---
 .../Conversion/SCFToAffine/scf-to-affine.mlir | 42 ++++++++++++++++++-
 2 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
index 488aae4c714c5..7113fbcd3c0e3 100644
--- a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -102,20 +102,40 @@ struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
                                  PatternRewriter &rewriter) const;
 };
 
+static bool areValidAffineMapOperands(AffineMap map, ValueRange operands,
+                                      Region *scope) {
+  assert(map.getNumInputs() == operands.size() &&
+         "expected one operand per affine map input");
+  return llvm::all_of(
+             operands.take_front(map.getNumDims()),
+             [&](Value value) { return affine::isValidDim(value, scope); }) &&
+         llvm::all_of(operands.drop_front(map.getNumDims()), [&](Value value) {
+           return affine::isValidSymbol(value, scope);
+         });
+}
+
 bool indexBoundsRaisable(scf::ForOp op) {
   Value lb = op.getLowerBound();
   Value ub = op.getUpperBound();
   IntegerAttr constAttr;
+  Region *scope = affine::getAffineScope(op);
+  if (!scope)
+    return false;
 
   // The asymmetry between lb and ub comes from the fact that the step
   // normalization (for non-constant (dynamic) steps) does not work with
   // multiple *lower* bounds (max).
-  bool lbOK = affine::isValidDim(lb) ||
-              (isa_and_present<affine::AffineMaxOp>(lb.getDefiningOp()) &&
-               matchPattern(op.getStep(), m_Constant(&constAttr)));
-  bool ubOK = affine::isValidDim(ub) ||
-              isa_and_present<affine::AffineMinOp>(ub.getDefiningOp());
-  bool stepOK = affine::isValidSymbol(op.getStep());
+  auto lbMaxOp = lb.getDefiningOp<affine::AffineMaxOp>();
+  bool lbOK = affine::isValidDim(lb, scope) ||
+              (lbMaxOp && matchPattern(op.getStep(), m_Constant(&constAttr)) &&
+               areValidAffineMapOperands(lbMaxOp.getAffineMap(),
+                                         lbMaxOp->getOperands(), scope));
+  auto ubMinOp = ub.getDefiningOp<affine::AffineMinOp>();
+  bool ubOK =
+      affine::isValidDim(ub, scope) ||
+      (ubMinOp && areValidAffineMapOperands(ubMinOp.getAffineMap(),
+                                            ubMinOp->getOperands(), scope));
+  bool stepOK = affine::isValidSymbol(op.getStep(), scope);
 
   return lbOK && ubOK && stepOK;
 }
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
index a489e3c3ce9ef..2a0ceddd1764f 100644
--- a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -1,7 +1,5 @@
 // RUN: mlir-opt --raise-scf-to-affine --split-input-file %s | FileCheck %s
 
-// XFAIL: mlir-expensive-checks
-
 // CHECK-LABEL: @constant_step
 // CHECK-SAME:  %[[ARR:.*]]: memref<?xi32>, %[[LB:.*]]: index, %[[UB:.*]]: index
 // CHECK:         affine.for %[[IV:.*]] = %[[LB]] to %[[UB]] step 3 {
@@ -155,6 +153,46 @@ func.func @constant_step_non_rectangular_nest() {
 
 // -----
 
+// CHECK-LABEL:   func.func @inner_not_raised_when_enclosing_loop_stuck
+// CHECK:           affine.for %{{.*}} = 0 to 4 {
+// CHECK:             memref.load
+// CHECK:             scf.for
+// CHECK:               affine.max
+// CHECK:               affine.min
+// CHECK:               scf.for
+// CHECK:                 func.call @some_func
+// CHECK-NOT:         affine.for
+
+#lbs = affine_map<(i)[K, N] -> (0, K - i)>
+#ubs = affine_map<(i)[K, N] -> (K, N - i)>
+
+func.func private @some_func(%i: index, %j: index)
+
+func.func @inner_not_raised_when_enclosing_loop_stuck(%mem: memref<?xindex>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+
+  %P = arith.constant 4 : index
+  %N = arith.constant 10 : index
+  %K = arith.constant 3 : index
+
+  scf.for %p = %c0 to %P step %c1 {
+    // %bad is neither top-level nor an affine quantity, so the middle loop can
+    // never be raised and %i never becomes a valid affine dimension.
+    %bad = memref.load %mem[%p] : memref<?xindex>
+    scf.for %i = %c0 to %bad step %c1 {
+      %lb = affine.max #lbs(%i)[%K, %N]
+      %ub = affine.min #ubs(%i)[%K, %N]
+      scf.for %j = %lb to %ub step %c1 {
+        func.call @some_func(%i, %j) : (index, index) -> ()
+      }
+    }
+  }
+  return
+}
+
+// -----
+
 // CHECK: #[[$UB_MAP:.+]] = affine_map<(d0)[s0] -> ((s0 + 5) floordiv s0, (-d0 + s0 + 98) floordiv s0)>
 // CHECK: #[[$IV_MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
 // CHECK-LABEL:   func.func @dynamic_step_non_rectangular_nest(