[Mlir-commits] [mlir] [mlir][SCF] Skip dynamic front-peeling for non-index loops (PR #218238)

[Federico Bruzzone]

https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/218238

>From b7f787d66fc488c0750883007225499788b0625e Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Sun, 23 Aug 2026 15:06:59 +0200
Subject: [PATCH 1/2] [mlir][SCF] Skip dynamic front-peeling for non-index
 loops

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../SCF/Transforms/LoopSpecialization.cpp     |  5 +++++
 .../Dialect/SCF/for-loop-peeling-front.mlir   | 21 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index ba5375bb42d94..27d0c32230f0b 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -236,6 +236,11 @@ LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
   if (lbInt && ubInt && stepInt && ceil(float(*ubInt - *lbInt) / *stepInt) <= 1)
     return failure();
 
+  // The peeling bound (%lb + %step) is computed with affine.apply, which
+  // accepts only index operands.
+  if ((!lbInt || !stepInt) && !forOp.getInductionVar().getType().isIndex())
+    return failure();
+
   AffineExpr lbSymbol, stepSymbol;
   bindSymbols(b.getContext(), lbSymbol, stepSymbol);
 
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
index 1737c6b6f6c5f..afd570c8b6f8c 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
@@ -218,3 +218,24 @@ func.func @no_peeling_front() -> i32 {
   }
   return %r : i32
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @non_index_loop_bounds(
+// CHECK-SAME:                                         %[[INIT:[^:]*]]: i32,
+// CHECK-SAME:                                         %[[LB:[^:]*]]: i32,
+// CHECK-SAME:                                         %[[N:[^:]*]]: i32) -> i32 {
+// CHECK:           %[[C2:.*]] = arith.constant 2 : i32
+// CHECK:           %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[N]] step %[[C2]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (i32) : i32 {
+// CHECK:             %[[ADD:.*]] = arith.addi %[[ACC]], %[[IV]] : i32
+// CHECK:             scf.yield %[[ADD]] : i32
+// CHECK:           }
+// CHECK:           return %[[RESULT]] : i32
+func.func @non_index_loop_bounds(%init: i32, %lb: i32, %n: i32) -> i32 {
+  %c2 = arith.constant 2 : i32
+  %r = scf.for %i = %lb to %n step %c2 iter_args(%a = %init) -> (i32) : i32 {
+    %t = arith.addi %a, %i : i32
+    scf.yield %t : i32
+  }
+  return %r : i32
+}

>From 90a6975dc8384f894668049477698c2fb32e80dc Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Mon, 24 Aug 2026 08:03:12 +0200
Subject: [PATCH 2/2] Address comments

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 27d0c32230f0b..bcc7ed3d3abfd 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -237,7 +237,9 @@ LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
     return failure();
 
   // The peeling bound (%lb + %step) is computed with affine.apply, which
-  // accepts only index operands.
+  // accepts only index operands. %ub does not feed into this bound, so only
+  // %lb and %step need to be constant to guarantee the affine.apply (see below)
+  // folds away before its (non-index) operand types matter.
   if ((!lbInt || !stepInt) && !forOp.getInductionVar().getType().isIndex())
     return failure();