[Mlir-commits] [mlir] [mlir][scf] refactor loop peeling logic to use arith operations (PR #235260)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 24 01:22:36 PDT 2026
https://github.com/StarOne01 created https://github.com/llvm/llvm-project/pull/235260
the signed split bound is now the arith chain instead of `affine.max` hence supporting peeling of signed non-index loops with dynamic bounds.
from #234435
claude helped me with tests and comments
>From bf386f6d707ab6da5cc0be005ac41e685e2fb44f Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 13:37:46 +0530
Subject: [PATCH] [mlir][scf] refactor loop peeling logic to use arith
operations and improve unsigned handling and modify tests
---
.../SCF/Transforms/LoopSpecialization.cpp | 95 +++++++------------
.../transform-op-peel-and-vectorize-conv.mlir | 5 +-
.../transform-op-peel-and-vectorize.mlir | 4 +-
mlir/test/Dialect/SCF/for-loop-peeling.mlir | 66 +++++++++----
.../SparseTensor/sparse_vector_peeled.mlir | 6 +-
5 files changed, 91 insertions(+), 85 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 46bf6883fcd89a..e232692ecc4e5a 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -132,57 +132,31 @@ static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
b.setInsertionPoint(forOp);
auto loc = forOp.getLoc();
- // Check whether we have a unsigned loop
- if (forOp.getUnsignedCmp()) {
- IntegerAttr lbAttr, ubAttr, stepAttr;
- if (matchPattern(lbVal, m_Constant(&lbAttr)) &&
- matchPattern(ubVal, m_Constant(&ubAttr)) &&
- matchPattern(stepVal, m_Constant(&stepAttr))) {
-
- // Get the constants as APInt at the loop’s own bit width
- APInt lb = lbAttr.getValue();
- APInt ub = ubAttr.getValue();
- APInt step = stepAttr.getValue();
-
- if (step.ule(1) || ub.ule(lb) || (ub - lb).urem(step).isZero())
- return failure();
- }
+ bool isUnsigned = forOp.getUnsignedCmp();
+ IntegerAttr lbAttr, ubAttr, stepAttr;
+ if (matchPattern(stepVal, m_Constant(&stepAttr))) {
+ APInt step = stepAttr.getValue();
+ if (isUnsigned ? step.ule(1) : step.sle(1))
+ return failure();
+ }
+ if (matchPattern(lbVal, m_Constant(&lbAttr)) &&
+ matchPattern(ubVal, m_Constant(&ubAttr))) {
- // split = ub - ((ub - min(lb, ub)) remui step). The min makes an empty loop
- // (ub <= lb) give a zero remainder, so split = ub and both loops stay
- // empty. The wrapping subtraction is exact for ub > lb, even above the
- // signed range.
- Value lbEff = b.createOrFold<arith::MinUIOp>(loc, lbVal, ubVal);
- Value diff = b.createOrFold<arith::SubIOp>(loc, ubVal, lbEff);
- Value rem = b.createOrFold<arith::RemUIOp>(loc, diff, stepVal);
- splitBound = b.createOrFold<arith::SubIOp>(loc, ubVal, rem);
+ APInt lb = lbAttr.getValue();
+ APInt ub = ubAttr.getValue();
- } else {
- auto lbInt = getConstantIntValue(lbVal);
- auto ubInt = getConstantIntValue(ubVal);
- auto stepInt = getConstantIntValue(stepVal);
- // No specialization necessary if step size is 1. Also bail out in case of
- // an invalid zero or negative step which might have happened during
- // folding.
- if (stepInt && *stepInt <= 1)
+ if (isUnsigned ? ub.ule(lb) : ub.sle(lb))
return failure();
- // No specialization necessary if step already divides upper bound evenly.
- // Fast path: lb, ub and step are constants.
- if (lbInt && ubInt && stepInt) {
- int64_t diff;
- // bail out if the subtraction overflows
- if (llvm::SubOverflow(*ubInt, *lbInt, diff) || diff % *stepInt == 0)
- return failure();
- }
-
- // Only the dynamic path computes the peeling bound with affine.apply, which
- // accepts only index operands.
- if ((!lbInt || !ubInt || !stepInt) &&
- !forOp.getInductionVar().getType().isIndex())
+ // lb < ub here, so ub - lb as an unsigned value is the exact distance, even
+ // for signed loops whose range exceeds the signed max. hence urem for both.
+ if (stepAttr && (ub - lb).urem(stepAttr.getValue()).isZero())
return failure();
+ }
- // Slow path: Examine the ops that define lb, ub and step.
+ // affine reasoning is signed and index-only,
+ // so use it only for signed index loops.
+ if (!isUnsigned && forOp.getInductionVar().getType().isIndex()) {
AffineExpr sym0, sym1, sym2;
bindSymbols(b.getContext(), sym0, sym1, sym2);
SmallVector<Value> operands{forOp.getLowerBound(), forOp.getUpperBound(),
@@ -192,20 +166,19 @@ static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
if (auto constExpr = dyn_cast<AffineConstantExpr>(map.getResult(0)))
if (constExpr.getValue() == 0)
return failure();
-
- // 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());
- 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);
}
+
+ // split = ub - ((ub - min(lb, ub)) remui step), with a signed or unsigned min
+ // to match the loop. For ub <= lb the min gives a zero remainder, so
+ // split = ub and both loops stay empty. For ub > lb the wrapping subtraction
+ // is the exact distance when read as unsigned, even beyond the signed range,
+ // which is why remui is used for signed loops as well.
+ Value lbEff = isUnsigned ? b.createOrFold<arith::MinUIOp>(loc, lbVal, ubVal)
+ : b.createOrFold<arith::MinSIOp>(loc, lbVal, ubVal);
+ Value diff = b.createOrFold<arith::SubIOp>(loc, ubVal, lbEff);
+ Value rem = b.createOrFold<arith::RemUIOp>(loc, diff, stepVal);
+ splitBound = b.createOrFold<arith::SubIOp>(loc, ubVal, rem);
+
// Create ForOp for partial iteration.
b.setInsertionPointAfter(forOp);
partialIteration = cast<ForOp>(b.clone(*forOp.getOperation()));
@@ -284,9 +257,9 @@ LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
// Check whether we have a unsigned loop
if (forOp.getUnsignedCmp()) {
IntegerAttr lbAttr, ubAttr, stepAttr;
- if (matchPattern(forOp.getLowerBound(), m_Constant(&lbAttr)) &&
- matchPattern(forOp.getUpperBound(), m_Constant(&ubAttr)) &&
- matchPattern(forOp.getStep(), m_Constant(&stepAttr))) {
+ if (matchPattern(lbVal, m_Constant(&lbAttr)) &&
+ matchPattern(ubVal, m_Constant(&ubAttr)) &&
+ matchPattern(stepVal, m_Constant(&stepAttr))) {
// Get the constants as APInt at the loop’s own bit width
APInt lb = lbAttr.getValue();
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 6e6de6ba139ab4..bee76a995290ef 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,8 +8,6 @@
// 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, 0)>
-
// CHECK-LABEL: func.func @conv(
// CHECK-DAG: %[[C_43:.*]] = arith.constant 43 : index
// CHECK-DAG: %[[C_48:.*]] = arith.constant 48 : index
@@ -20,7 +18,8 @@ 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.max #[[$MAP]](){{\[}}%[[VSCALE_X_4]]]
+// CHECK: %[[REM:.*]] = arith.remui %[[C_48]], %[[VSCALE_X_4]] : index
+// CHECK-NEXT: %[[UB_DEPTH_LOOP:.*]] = arith.subi %[[C_48]], %[[REM]] : index
// 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 38c1bb59b4dfb4..2818b29d3f45c5 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,6 @@ func.func @matmul(%A: tensor<1024x512xf32>,
%B: tensor<512x2000xf32>,
%C: tensor<1024x2000xf32>) -> tensor<1024x2000xf32> {
-// 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 +22,8 @@ func.func @matmul(%A: tensor<1024x512xf32>,
// Main loop after vectorisation (without masking)
-// CHECK: %[[UB_MAIN:.*]] = affine.max #[[MAP]]()[%[[STEP]]]
+// CHECK: %[[REM:.*]] = arith.remui %[[C2000]], %[[STEP]] : index
+// CHECK-NEXT: %[[UB_MAIN:.*]] = arith.subi %[[C2000]], %[[REM]] : index
// 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.mlir b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
index 59febdaa55fd7e..35b1ea6cd174f7 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
@@ -2,12 +2,14 @@
// 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, 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.max #[[MAP0]]()[%[[LB]], %[[UB]], %[[STEP]]]
+// CHECK: %[[LB_EFF:.*]] = arith.minsi %[[LB]], %[[UB]] : index
+// CHECK: %[[DIFF:.*]] = arith.subi %[[UB]], %[[LB_EFF]] : index
+// CHECK: %[[SPLIT_REM:.*]] = arith.remui %[[DIFF]], %[[STEP]] : index
+// CHECK: %[[NEW_UB:.*]] = arith.subi %[[UB]], %[[SPLIT_REM]] : index
// 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 +134,6 @@ func.func @fully_static_bounds_integers() -> i32 {
// -----
-// 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 +141,10 @@ 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.max #[[MAP0]]()[%[[UB]]]
+// CHECK: %[[LB_EFF:.*]] = arith.minsi %[[UB]], %[[C0]] : index
+// CHECK: %[[DIFF:.*]] = arith.subi %[[UB]], %[[LB_EFF]] : index
+// CHECK: %[[SPLIT_REM:.*]] = arith.remui %[[DIFF]], %[[C4]] : index
+// CHECK: %[[NEW_UB:.*]] = arith.subi %[[UB]], %[[SPLIT_REM]] : index
// 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 +175,16 @@ func.func @dynamic_upper_bound(%ub : index) -> i32 {
// -----
-// 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.max #[[MAP0]]()[%[[UB]]]
+// CHECK: %[[LB_EFF:.*]] = arith.minsi %[[UB]], %[[C0]] : index
+// CHECK: %[[DIFF:.*]] = arith.subi %[[UB]], %[[LB_EFF]] : index
+// CHECK: %[[SPLIT_REM:.*]] = arith.remui %[[DIFF]], %[[C4]] : index
+// CHECK: %[[NEW_UB:.*]] = arith.subi %[[UB]], %[[SPLIT_REM]] : index
// CHECK: scf.for %[[IV:.*]] = %[[C0]] to %[[NEW_UB]] step %[[C4]] {
// CHECK: %[[LOAD:.*]] = memref.load %[[MEMREF]][]
// CHECK: %[[ADD:.*]] = arith.addi %[[LOAD]], %[[C4_I32]] : i32
@@ -360,12 +366,9 @@ func.func @regression(%arg0: memref<i64>, %arg1: index) {
// Regression test: Make sure that we do not crash.
// CHECK-LABEL: func @zero_step(
-// CHECK: %[[c0:.*]] = arith.constant 0
-// CHECK: %[[c1:.*]] = arith.constant 1
-// CHECK: %[[bound:.*]] = affine.max
-// CHECK: scf.for %{{.*}} = %[[c0]] to %[[bound]] step %[[c0]]
+// CHECK: scf.for
// CHECK: arith.index_cast
-// CHECK: scf.for %{{.*}} = %[[bound]] to %[[c1]] step %[[c0]]
+// CHECK: scf.for
// CHECK: arith.index_cast
func.func @zero_step(%arg0: memref<i64>) {
%c0 = arith.constant 0 : index
@@ -383,12 +386,14 @@ func.func @zero_step(%arg0: memref<i64>) {
// CHECK-LABEL: func.func @non_index_loop_bounds(
// CHECK-SAME: %[[INIT:.*]]: i32,
// CHECK-SAME: %[[N:.*]]: i32) -> i32 {
-// CHECK: %[[C0:.*]] = arith.constant 0 : i32
-// CHECK: %[[C2:.*]] = arith.constant 2 : i32
-// CHECK: %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[N]] step %[[C2]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (i32) : i32 {
-// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[IV]] : i32
-// CHECK: scf.yield %[[ADD]] : i32
-// CHECK: }
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
+// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i32
+// CHECK: %[[LB_EFF:.*]] = arith.minsi %[[N]], %[[C0]] : i32
+// CHECK: %[[DIFF:.*]] = arith.subi %[[N]], %[[LB_EFF]] : i32
+// CHECK: %[[SPLIT_REM:.*]] = arith.remui %[[DIFF]], %[[C2]] : i32
+// CHECK: %[[NEW_UB:.*]] = arith.subi %[[N]], %[[SPLIT_REM]] : i32
+// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[NEW_UB]] step %[[C2]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (i32) : i32 {
+// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[NEW_UB]] to %[[N]] step %[[C2]] iter_args(%[[ACC2:.*]] = %[[FIRST]]) -> (i32) : i32 {
// CHECK: return %[[RESULT]] : i32
func.func @non_index_loop_bounds(%init: i32, %n: i32) -> i32 {
%c0 = arith.constant 0 : i32
@@ -565,3 +570,30 @@ func.func @signed_empty_loop_not_executed() {
}
return
}
+
+// -----
+
+// Signed i32 loop whose range (INT_MIN to INT_MAX - 4) exceeds the signed max.
+// ub - lb = 4294967291 only makes sense as an unsigned value, so remui is
+// needed for signed loops too. remainder 3, so split = 2147483640. The
+// partial loop is exactly one iteration, so canonicalization inlines it.
+
+func.func private @effect(i32)
+
+// CHECK-LABEL: func @signed_range_above_signed_max
+// CHECK-DAG: %[[LB:.*]] = arith.constant -2147483648 : i32
+// CHECK-DAG: %[[STEP:.*]] = arith.constant 4 : i32
+// CHECK-DAG: %[[SPLIT:.*]] = arith.constant 2147483640 : i32
+// CHECK: scf.for %{{.*}} = %[[LB]] to %[[SPLIT]] step %[[STEP]] : i32 {
+// CHECK: call @effect(%{{.*}}) : (i32) -> ()
+// CHECK: }
+// CHECK: call @effect(%[[SPLIT]]) : (i32) -> ()
+func.func @signed_range_above_signed_max() {
+ %lb = arith.constant -2147483648 : i32
+ %ub = arith.constant 2147483643 : i32
+ %step = arith.constant 4 : i32
+ scf.for %iv = %lb to %ub step %step : i32 {
+ func.call @effect(%iv) : (i32) -> ()
+ }
+ return
+}
diff --git a/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
index ed0ede6da5669e..62491f6d746368 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
@@ -17,7 +17,6 @@
doc = "x(i) = a(i) * b(i)"
}
-// 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 +29,10 @@
// 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.max #[[$map0]]()[%[[q]], %[[s]]]
+// CHECK: %[[lbeff:.*]] = arith.minsi %[[q]], %[[s]] : index
+// CHECK: %[[diff:.*]] = arith.subi %[[s]], %[[lbeff]] : index
+// CHECK: %[[rem:.*]] = arith.remui %[[diff]], %[[c16]] : index
+// CHECK: %[[boundary:.*]] = arith.subi %[[s]], %[[rem]] : index
// 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