[Mlir-commits] [mlir] [mlir][scf] Do not peel an iteration out of an empty loop (PR #223257)
Samarth Narang
llvmlistbot at llvm.org
Sun Sep 13 11:07:38 PDT 2026
https://github.com/snarang181 updated https://github.com/llvm/llvm-project/pull/223257
>From cb944266884afb8c2ed63ae7c8a9a65c4429e277 Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Sun, 13 Sep 2026 11:39:39 +0000
Subject: [PATCH] [mlir][scf] Do not peel an iteration out of an empty loop
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.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
.../SCF/Transforms/LoopSpecialization.cpp | 30 ++++++++++++-------
.../transform-op-peel-and-vectorize-conv.mlir | 4 +--
.../transform-op-peel-and-vectorize.mlir | 4 +--
.../Dialect/SCF/for-loop-peeling-front.mlir | 11 ++++---
mlir/test/Dialect/SCF/for-loop-peeling.mlir | 18 +++++------
.../SparseTensor/sparse_vector_peeled.mlir | 4 +--
6 files changed, 42 insertions(+), 29 deletions(-)
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/Linalg/transform-op-peel-and-vectorize-conv.mlir b/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize-conv.mlir
index 9805372233a64..6e6de6ba139ab 100644
--- a/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize-conv.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize-conv.mlir
@@ -8,7 +8,7 @@
// masks in the main loop.
func.func @conv(%arg0: tensor<1x1080x1962x48xi32>, %arg1: tensor<1x43x48xi32>) -> tensor<1x1080x1920x48xi32> {
-// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (-(48 mod s0) + 48)>
+// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (-(48 mod s0) + 48, 0)>
// CHECK-LABEL: func.func @conv(
// CHECK-DAG: %[[C_43:.*]] = arith.constant 43 : index
@@ -20,7 +20,7 @@ func.func @conv(%arg0: tensor<1x1080x1962x48xi32>, %arg1: tensor<1x43x48xi32>) -
// CHECK: %[[VSCALE_X_4:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
// Loop over the channel/depth dim - the main part after vectorisation (vectorized, no masking)
-// CHECK: %[[UB_DEPTH_LOOP:.*]] = affine.apply #[[$MAP]](){{\[}}%[[VSCALE_X_4]]]
+// CHECK: %[[UB_DEPTH_LOOP:.*]] = affine.max #[[$MAP]](){{\[}}%[[VSCALE_X_4]]]
// CHECK-NEXT: %[[VAL_21:.*]] = scf.for {{.*}} to %[[UB_DEPTH_LOOP]] step %[[VSCALE_X_4]]
// Loop over the Filter width dim
// CHECK: scf.for %{{.*}} = %[[C0]] to %[[C_43]] step %[[C1]] {{.*}} -> (tensor<1x1x4x?xi32>) {
diff --git a/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize.mlir b/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize.mlir
index 05a032b1ece06..38c1bb59b4dfb 100644
--- a/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize.mlir
@@ -9,7 +9,7 @@ func.func @matmul(%A: tensor<1024x512xf32>,
%B: tensor<512x2000xf32>,
%C: tensor<1024x2000xf32>) -> tensor<1024x2000xf32> {
-// CHECK: #[[MAP:.*]] = affine_map<()[s0] -> (-(2000 mod s0) + 2000)>
+// CHECK: #[[MAP:.*]] = affine_map<()[s0] -> (-(2000 mod s0) + 2000, 0)>
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C2000:.*]] = arith.constant 2000 : index
// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
@@ -23,7 +23,7 @@ func.func @matmul(%A: tensor<1024x512xf32>,
// Main loop after vectorisation (without masking)
-// CHECK: %[[UB_MAIN:.*]] = affine.apply #[[MAP]]()[%[[STEP]]]
+// CHECK: %[[UB_MAIN:.*]] = affine.max #[[MAP]]()[%[[STEP]]]
// CHECK: scf.for {{.*}} %[[C0]] to %[[UB_MAIN]] step %[[STEP]] {{.*}} -> (tensor<1024x2000xf32>) {
// CHECK: scf.for %arg7 = %[[C0]] to %[[C512]] step %[[C1]] {{.*}} -> (tensor<1024x2000xf32>) {
// CHECK-NOT: vector.mask
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
diff --git a/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
index 35fd7c33e4cfe..ed0ede6da5669 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
@@ -17,7 +17,7 @@
doc = "x(i) = a(i) * b(i)"
}
-// CHECK-DAG: #[[$map0:.*]] = affine_map<()[s0, s1] -> (s0 + ((-s0 + s1) floordiv 16) * 16)>
+// CHECK-DAG: #[[$map0:.*]] = affine_map<()[s0, s1] -> (s0 + ((-s0 + s1) floordiv 16) * 16, s0)>
// CHECK-DAG: #[[$map1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
// CHECK-LABEL: func @mul_s
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
@@ -30,7 +30,7 @@
// CHECK: %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32>
// CHECK: %[[b:.*]] = arith.extui %[[r]] : i32 to i64
// CHECK: %[[s:.*]] = arith.index_cast %[[b]] : i64 to index
-// CHECK: %[[boundary:.*]] = affine.apply #[[$map0]]()[%[[q]], %[[s]]]
+// CHECK: %[[boundary:.*]] = affine.max #[[$map0]]()[%[[q]], %[[s]]]
// CHECK: scf.for %[[i:.*]] = %[[q]] to %[[boundary]] step %[[c16]] {
// CHECK: %[[li:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xi32>, vector<16xi32>
// CHECK: %[[zi:.*]] = arith.extui %[[li]] : vector<16xi32> to vector<16xi64>
More information about the Mlir-commits
mailing list