[Mlir-commits] [mlir] [mlir][scf] Do not peel an iteration out of an empty loop (PR #223257)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Sep 13 10:24:33 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Samarth Narang (snarang181)
<details>
<summary>Changes</summary>
Both peeling directions bail out only when every bound is a bare constant, so a loop with any dynamic bound is peeled unconditionally, including one whose runtime trip count is zero. The peeled iteration then executes a body the source never executes.
Rather than extending the constant-only guards, make the split bound correct for any runtime state: clamp the back-peeling bound to at least %lb, and take the front-peeling bound as min(%lb + %step, %ub). Both are no-ops whenever the loop has at least one iteration.
Fixes #<!-- -->223232.
---
Full diff: https://github.com/llvm/llvm-project/pull/223257.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp (+20-10)
- (modified) mlir/test/Dialect/SCF/for-loop-peeling-front.mlir (+7-4)
- (modified) mlir/test/Dialect/SCF/for-loop-peeling.mlir (+9-9)
``````````diff
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 2f15c0414b660..a0b5ae9b66091 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -154,14 +154,17 @@ static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
if (constExpr.getValue() == 0)
return failure();
- // New upper bound: %ub - (%ub - %lb) mod %step
- auto modMap = AffineMap::get(0, 3, {sym1 - ((sym1 - sym0) % sym2)});
+ // New upper bound: max(%lb, %ub - (%ub - %lb) mod %step). Affine `mod` is
+ // non-negative, so without the clamp an empty loop (%lb >= %ub) gets a split
+ // bound below %lb and its partial iteration executes.
+ auto modMap = AffineMap::get(0, 3, {sym1 - ((sym1 - sym0) % sym2), sym0},
+ b.getContext());
b.setInsertionPoint(forOp);
auto loc = forOp.getLoc();
- splitBound = b.createOrFold<AffineApplyOp>(loc, modMap,
- ValueRange{forOp.getLowerBound(),
- forOp.getUpperBound(),
- forOp.getStep()});
+ splitBound = b.createOrFold<AffineMaxOp>(loc, modMap,
+ ValueRange{forOp.getLowerBound(),
+ forOp.getUpperBound(),
+ forOp.getStep()});
if (splitBound.getType() != forOp.getLowerBound().getType())
splitBound = b.createOrFold<arith::IndexCastOp>(
loc, forOp.getLowerBound().getType(), splitBound);
@@ -246,12 +249,19 @@ LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
AffineExpr lbSymbol, stepSymbol;
bindSymbols(b.getContext(), lbSymbol, stepSymbol);
- // New lower bound for main loop: %lb + %step
- auto ubMap = AffineMap::get(0, 2, {lbSymbol + stepSymbol});
+ // New lower bound for main loop: min(%lb + %step, %ub). Without the minimum
+ // the peeled first iteration spans [%lb, %lb + %step) and executes even when
+ // the source loop is empty (%lb >= %ub).
+ AffineExpr ubSymbol;
+ bindSymbols(b.getContext(), lbSymbol, stepSymbol, ubSymbol);
+ auto ubMap =
+ AffineMap::get(0, 3, {lbSymbol + stepSymbol, ubSymbol}, b.getContext());
b.setInsertionPoint(forOp);
auto loc = forOp.getLoc();
- Value splitBound = b.createOrFold<AffineApplyOp>(
- loc, ubMap, ValueRange{forOp.getLowerBound(), forOp.getStep()});
+ Value splitBound = b.createOrFold<AffineMinOp>(
+ loc, ubMap,
+ ValueRange{forOp.getLowerBound(), forOp.getStep(),
+ forOp.getUpperBound()});
if (splitBound.getType() != forOp.getUpperBound().getType())
splitBound = b.createOrFold<arith::IndexCastOp>(
loc, forOp.getUpperBound().getType(), splitBound);
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
index afd570c8b6f8c..ec91e4def6484 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
@@ -85,14 +85,17 @@ func.func @static_two_iterations_ub_used_in_loop(%arg0: memref<1xi32>) -> i32 {
// CHECK-SAME: %[[UB:.*]]: index, %[[MEMREF:.*]]: memref<i32>
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-// CHECK: scf.for %[[IV:.*]] = %[[C0]] to %[[C4]] step %[[C4]] {
+// The peeled first iteration is bounded by %ub as well as by %lb + %step, so
+// it stays empty when the source loop has no iterations.
+// CHECK: %[[SPLIT:.*]] = affine.min {{.*}}[%[[C0]], %[[C4]], %[[UB]]]
+// CHECK: scf.for %[[IV:.*]] = %[[C0]] to %[[SPLIT]] step %[[C4]] {
// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[UB]], %[[IV]])[%[[C4]]]
// CHECK: %[[LOAD:.*]] = memref.load %[[MEMREF]][]
// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]]
// CHECK: %[[ADD:.*]] = arith.addi %[[LOAD]], %[[CAST]] : i32
// CHECK: memref.store %[[ADD]], %[[MEMREF]]
// CHECK: }
-// CHECK: scf.for %[[IV2:.*]] = %[[C4]] to %[[UB]] step %[[C4]] {
+// CHECK: scf.for %[[IV2:.*]] = %[[SPLIT]] to %[[UB]] step %[[C4]] {
// CHECK: %[[REM:.*]] = affine.min #[[MAP]](%[[UB]], %[[IV2]])[%[[C4]]]
// CHECK: %[[LOAD2:.*]] = memref.load %[[MEMREF]][]
// CHECK: %[[CAST2:.*]] = arith.index_cast %[[REM]]
@@ -117,12 +120,12 @@ func.func @no_loop_results(%ub : index, %d : memref<i32>) {
// -----
-// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1] -> (s0 + s1)>
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s0 + s1, s2)>
// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
// CHECK: func @fully_dynamic_bounds(
// CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
// CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32
-// CHECK: %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[LB]], %[[STEP]]]
+// CHECK: %[[NEW_UB:.*]] = affine.min #[[MAP0]]()[%[[LB]], %[[STEP]], %[[UB]]]
// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]]
// CHECK-SAME: step %[[STEP]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {
// CHECK: %[[MIN:.*]] = affine.min #[[MAP1]](%[[UB]], %[[IV]])[%[[STEP]]]
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling.mlir b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
index d1026f5090cb0..86d40f176e63e 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
@@ -2,12 +2,12 @@
// RUN: mlir-opt %s -scf-for-loop-peeling=skip-partial=false -canonicalize -split-input-file -verify-diagnostics | FileCheck %s -check-prefix=CHECK-NO-SKIP
-// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (-s0 + s1) mod s2)>
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (-s0 + s1) mod s2, s0)>
// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
// CHECK: func @fully_dynamic_bounds(
// CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
// CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32
-// CHECK: %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[LB]], %[[UB]], %[[STEP]]]
+// CHECK: %[[NEW_UB:.*]] = affine.max #[[MAP0]]()[%[[LB]], %[[UB]], %[[STEP]]]
// CHECK: %[[LOOP:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]]
// CHECK-SAME: step %[[STEP]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {
// CHECK: %[[CAST:.*]] = arith.index_cast %[[STEP]] : index to i32
@@ -132,7 +132,7 @@ func.func @fully_static_bounds_integers() -> i32 {
// -----
-// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> ((s0 floordiv 4) * 4)>
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> ((s0 floordiv 4) * 4, 0)>
// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
// CHECK: func @dynamic_upper_bound(
// CHECK-SAME: %[[UB:.*]]: index
@@ -140,7 +140,7 @@ func.func @fully_static_bounds_integers() -> i32 {
// CHECK-DAG: %[[C4_I32:.*]] = arith.constant 4 : i32
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
-// CHECK: %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[UB]]]
+// CHECK: %[[NEW_UB:.*]] = affine.max #[[MAP0]]()[%[[UB]]]
// CHECK: %[[LOOP:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[NEW_UB]]
// CHECK-SAME: step %[[C4]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {
// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[C4_I32]] : i32
@@ -171,14 +171,14 @@ func.func @dynamic_upper_bound(%ub : index) -> i32 {
// -----
-// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> ((s0 floordiv 4) * 4)>
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> ((s0 floordiv 4) * 4, 0)>
// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
// CHECK: func @no_loop_results(
// CHECK-SAME: %[[UB:.*]]: index, %[[MEMREF:.*]]: memref<i32>
// CHECK-DAG: %[[C4_I32:.*]] = arith.constant 4 : i32
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
-// CHECK: %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[UB]]]
+// CHECK: %[[NEW_UB:.*]] = affine.max #[[MAP0]]()[%[[UB]]]
// CHECK: scf.for %[[IV:.*]] = %[[C0]] to %[[NEW_UB]] step %[[C4]] {
// CHECK: %[[LOAD:.*]] = memref.load %[[MEMREF]][]
// CHECK: %[[ADD:.*]] = arith.addi %[[LOAD]], %[[C4_I32]] : i32
@@ -362,10 +362,10 @@ func.func @regression(%arg0: memref<i64>, %arg1: index) {
// CHECK-LABEL: func @zero_step(
// CHECK: %[[c0:.*]] = arith.constant 0
// CHECK: %[[c1:.*]] = arith.constant 1
-// CHECK: %[[poison:.*]] = ub.poison
-// CHECK: scf.for %{{.*}} = %[[c0]] to %[[poison]] step %[[c0]]
+// CHECK: %[[bound:.*]] = affine.max
+// CHECK: scf.for %{{.*}} = %[[c0]] to %[[bound]] step %[[c0]]
// CHECK: arith.index_cast
-// CHECK: scf.for %{{.*}} = %[[poison]] to %[[c1]] step %[[c0]]
+// CHECK: scf.for %{{.*}} = %[[bound]] to %[[c1]] step %[[c0]]
// CHECK: arith.index_cast
func.func @zero_step(%arg0: memref<i64>) {
%c0 = arith.constant 0 : index
``````````
</details>
https://github.com/llvm/llvm-project/pull/223257
More information about the Mlir-commits
mailing list