[Mlir-commits] [mlir] [mlir][scf] Clamp normalized loop sizes where they are combined (PR #223256)
Samarth Narang
llvmlistbot at llvm.org
Mon Oct 5 17:54:49 PDT 2026
https://github.com/snarang181 updated https://github.com/llvm/llvm-project/pull/223256
>From b71647e5f55cb654d3db5bf1956134baf97ca280 Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Sun, 13 Sep 2026 11:57:08 +0000
Subject: [PATCH 1/3] [mlir][scf] Clamp normalized loop sizes where they are
combined
emitNormalizedLoopBounds returns ceilDiv(ub - lb, step), which is negative for
an empty loop. A single negative size is harmless -- the loop is empty either
way -- but coalesceLoops and collapseParallelLoops multiply sizes across
dimensions, and two negatives cancel into a positive extent. A nest of two empty
loops then executes its body four times. The negative value is also not a legal
affine.delinearize_index basis.
Clamp at the points where sizes are combined rather than inside
emitNormalizedLoopBounds, so loops that are merely normalized keep their bound
unchanged.
Fixes #223231.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 26 ++++++++++++++++---
mlir/test/Dialect/Affine/loop-coalescing.mlir | 9 ++++---
.../Dialect/SCF/transform-op-coalesce.mlir | 20 +++++++++-----
3 files changed, 41 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 7ba0a796c8631..92a969b10091b 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -989,6 +989,21 @@ delinearizeInductionVariable(RewriterBase &rewriter, Location loc,
return {delinearizedIvs, preservedUsers};
}
+/// Clamp a normalized loop size at zero. `emitNormalizedLoopBounds` returns
+/// `ceilDiv(ub - lb, step)`, which is negative for an empty loop: combining two
+/// such sizes cancels the signs into a positive extent, and a negative value is
+/// not a legal `affine.delinearize_index` basis.
+static Value clampSizeToNonNegative(RewriterBase &rewriter, Location loc,
+ Value size) {
+ if (auto cst = getConstantIntValue(size))
+ return *cst >= 0 ? size
+ : getValueOrCreateConstantIntOp(rewriter, loc,
+ rewriter.getIndexAttr(0));
+ Value zero = arith::ConstantOp::create(rewriter, loc,
+ rewriter.getZeroAttr(size.getType()));
+ return arith::MaxSIOp::create(rewriter, loc, size, zero);
+}
+
LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
MutableArrayRef<scf::ForOp> loops) {
if (loops.size() < 2)
@@ -1035,8 +1050,9 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(outermost);
Location loc = outermost.getLoc();
- SmallVector<Value> upperBounds = llvm::map_to_vector(
- loops, [](auto loop) { return loop.getUpperBound(); });
+ SmallVector<Value> upperBounds = llvm::map_to_vector(loops, [&](auto loop) {
+ return clampSizeToNonNegative(rewriter, loc, loop.getUpperBound());
+ });
Value upperBound = getProductOfIntsOrIndexes(rewriter, loc, upperBounds);
outermost.setUpperBound(upperBound);
@@ -1188,8 +1204,10 @@ void mlir::collapseParallelLoops(
Value ub = loops.getUpperBound()[i];
Value step = loops.getStep()[i];
auto newLoopRange = emitNormalizedLoopBounds(rewriter, loc, lb, ub, step);
- normalizedUpperBounds.push_back(getValueOrCreateConstantIntOp(
- rewriter, loops.getLoc(), newLoopRange.size));
+ normalizedUpperBounds.push_back(clampSizeToNonNegative(
+ rewriter, loops.getLoc(),
+ getValueOrCreateConstantIntOp(rewriter, loops.getLoc(),
+ newLoopRange.size)));
rewriter.setInsertionPointToStart(loops.getBody());
denormalizeInductionVariable(rewriter, loc, loops.getInductionVars()[i], lb,
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index b8a68b9c5c1f9..b33a3b1efc01c 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -226,9 +226,12 @@ func.func @parametric(%lb1 : index, %ub1 : index, %step1 : index,
// CHECK: %[[c1:.+]] = arith.constant 1
// CHECK: %[[normalized_j:.*]] = affine.apply
// CHECK-SAME: affine_map<()[s0, s1, s2] -> ((-s0 + s1) ceildiv s2)>()[%[[orig_lb2]], %[[orig_ub2]], %[[orig_step2]]]
+ // Each normalized size is clamped at zero before the sizes are combined, so
+ // two empty dimensions cannot cancel into a positive flattened extent.
+ // CHECK: %[[clamped_i:.+]] = arith.maxsi %[[normalized_i]]
+ // CHECK: %[[clamped_j:.+]] = arith.maxsi %[[normalized_j]]
// CHECK: %[[range:.+]] = affine.apply
- // CHECK-SAME: affine_map<()[s0, s1, s2, s3, s4, s5] -> (((-s0 + s1) ceildiv s2) * ((-s3 + s4) ceildiv s5))>()
- // CHECK-SAME: [%[[orig_lb1]], %[[orig_ub1]], %[[orig_step1]], %[[orig_lb2]], %[[orig_ub2]], %[[orig_step2]]]
+ // CHECK-SAME: affine_map<()[s0, s1] -> (s0 * s1)>()[%[[clamped_i]], %[[clamped_j]]]
// Check that the outer loop is updated.
// CHECK: scf.for %[[i:.*]] = %[[c0]] to %[[range]] step %[[c1]]
@@ -237,7 +240,7 @@ func.func @parametric(%lb1 : index, %ub1 : index, %step1 : index,
// CHECK-NOT: scf.for
scf.for %j = %lb2 to %ub2 step %step2 {
// Remapping of the induction variables.
- // CHECK: %[[delinearize:.+]]:2 = affine.delinearize_index %[[i]] into (%[[normalized_i]], %[[normalized_j]])
+ // CHECK: %[[delinearize:.+]]:2 = affine.delinearize_index %[[i]] into (%[[clamped_i]], %[[clamped_j]])
// CHECK: %[[orig_j:.*]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
// CHECK-SAME: (%[[delinearize]]#1)[%[[orig_lb2]], %[[orig_step2]]]
// CHECK: %[[orig_i:.*]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
diff --git a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
index b487b62f0e3fe..bbf7ff986bba8 100644
--- a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
+++ b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
@@ -141,11 +141,14 @@ module attributes {transform.with_named_sequence} {
// CHECK-SAME: affine_map<()[s0, s1, s2] -> ((-s0 + s1) ceildiv s2)>()[%[[LB1]], %[[UB1]], %[[STEP1]]]
// CHECK: %[[NITERS2:.+]] = affine.apply
// CHECK-SAME: affine_map<()[s0, s1, s2] -> ((-s0 + s1) ceildiv s2)>()[%[[LB2]], %[[UB2]], %[[STEP2]]]
-// CHECK: %[[NEWUB:.+]] = affine.apply affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8] ->
-// CHECK-SAME: ((((-s0 + s1) ceildiv s2) * ((-s3 + s4) ceildiv s5)) * ((-s6 + s7) ceildiv s8))
-// CHECK-SAME: [%[[LB0]], %[[UB0]], %[[STEP0]], %[[LB1]], %[[UB1]], %[[STEP1]], %[[LB2]], %[[UB2]], %[[STEP2]]]
+// Each size is clamped at zero before the sizes are combined, so empty
+// dimensions cannot cancel into a positive flattened extent.
+// CHECK: %[[CL0:.+]] = arith.maxsi %[[NITERS0]]
+// CHECK: %[[CL1:.+]] = arith.maxsi %[[NITERS1]]
+// CHECK: %[[CL2:.+]] = arith.maxsi %[[NITERS2]]
+// CHECK: %[[NEWUB:.+]] = affine.apply
// CHECK: %[[RESULT:.+]] = scf.for %[[IV:[a-zA-Z0-9]+]] = %[[C0]] to %[[NEWUB]] step %[[C1]] iter_args(%[[ITER_ARG:.+]] = %[[ARG0]])
-// CHECK: %[[DELINEARIZE:.+]]:3 = affine.delinearize_index %[[IV]] into (%[[NITERS0]], %[[NITERS1]], %[[NITERS2]])
+// CHECK: %[[DELINEARIZE:.+]]:3 = affine.delinearize_index %[[IV]] into (%[[CL0]], %[[CL1]], %[[CL2]])
// CHECK-DAG: %[[K:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#2)[%[[LB2]], %[[STEP2]]]
// CHECK-DAG: %[[J:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#1)[%[[LB1]], %[[STEP1]]]
// CHECK-DAG: %[[I:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#0)[%[[LB0]], %[[STEP0]]]
@@ -324,9 +327,11 @@ module attributes {transform.with_named_sequence} {
// CHECK-SAME: %[[ARG2:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK: %[[UB:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%[[ARG1]], %[[ARG2]]]
+// CHECK: %[[CL1:.+]] = arith.maxsi %[[ARG1]]
+// CHECK: %[[CL2:.+]] = arith.maxsi %[[ARG2]]
+// CHECK: %[[UB:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%[[CL1]], %[[CL2]]]
// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[UB]] step %[[C1]]
-// CHECK: %[[DELINEARIZE:.+]]:2 = affine.delinearize_index %[[IV]] into (%[[ARG1]], %[[ARG2]])
+// CHECK: %[[DELINEARIZE:.+]]:2 = affine.delinearize_index %[[IV]] into (%[[CL1]], %[[CL2]])
// CHECK: "some_use"(%{{[a-zA-Z0-9]+}}, %[[C0]], %[[C0]], %[[DELINEARIZE]]#0, %[[C0]], %[[DELINEARIZE]]#1)
// -----
@@ -363,7 +368,8 @@ module attributes {transform.with_named_sequence} {
// CHECK-SAME: , %[[ARG1:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[ARG1]] step %[[C1]]
+// CHECK: %[[CL:.+]] = arith.maxsi %[[ARG1]]
+// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[CL]] step %[[C1]]
// CHECK: "some_use"(%{{[a-zA-Z0-9]+}}, %[[C0]], %[[C0]], %[[IV]])
// -----
>From e40dff4c7c9ab8e73dff7c9864aa7029c8665946 Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Thu, 24 Sep 2026 11:56:11 +0000
Subject: [PATCH 2/3] Address review: clamp the delinearization basis at one
Clamping only the combined upper bound at zero is enough to give the flattened
loop the right trip count, but it leaves the values used to recover the
original induction variables (the affine.delinearize_index basis in
coalesceLoops, the remsi/divsi divisors in collapseParallelLoops) at zero for
an empty dimension. A zero basis is undefined behaviour dynamically and a
verifier error once it is a constant. Clamp those at one instead: when any size
is non-positive the loop body never executes, and otherwise the clamp is the
identity. Use the size's own type for the clamp constant.
Add a statically empty nest to transform-op-coalesce.mlir, which folds to a
zero-trip loop with basis (1, 1), and a dynamic collapseParallelLoops case.
---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 44 ++++++++-------
mlir/test/Dialect/Affine/loop-coalescing.mlir | 11 ++--
.../Dialect/SCF/transform-op-coalesce.mlir | 54 +++++++++++++++----
.../single-parallel-loop-collapsing.mlir | 23 ++++++++
4 files changed, 99 insertions(+), 33 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 92a969b10091b..1d8a9641d11fc 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -989,19 +989,17 @@ delinearizeInductionVariable(RewriterBase &rewriter, Location loc,
return {delinearizedIvs, preservedUsers};
}
-/// Clamp a normalized loop size at zero. `emitNormalizedLoopBounds` returns
-/// `ceilDiv(ub - lb, step)`, which is negative for an empty loop: combining two
-/// such sizes cancels the signs into a positive extent, and a negative value is
-/// not a legal `affine.delinearize_index` basis.
-static Value clampSizeToNonNegative(RewriterBase &rewriter, Location loc,
- Value size) {
+/// Clamp a normalized loop size, `ceilDiv(ub - lb, step)`, which is negative
+/// for an empty loop. Sizes multiplied into an upper bound are clamped at zero;
+/// sizes used as a delinearization basis or divisor at one.
+static Value clampSize(RewriterBase &rewriter, Location loc, Value size,
+ int64_t minValue) {
+ TypedAttr minAttr = rewriter.getIntegerAttr(size.getType(), minValue);
if (auto cst = getConstantIntValue(size))
- return *cst >= 0 ? size
- : getValueOrCreateConstantIntOp(rewriter, loc,
- rewriter.getIndexAttr(0));
- Value zero = arith::ConstantOp::create(rewriter, loc,
- rewriter.getZeroAttr(size.getType()));
- return arith::MaxSIOp::create(rewriter, loc, size, zero);
+ return *cst >= minValue ? size
+ : arith::ConstantOp::create(rewriter, loc, minAttr);
+ Value min = arith::ConstantOp::create(rewriter, loc, minAttr);
+ return arith::MaxSIOp::create(rewriter, loc, size, min);
}
LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
@@ -1051,7 +1049,10 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
rewriter.setInsertionPoint(outermost);
Location loc = outermost.getLoc();
SmallVector<Value> upperBounds = llvm::map_to_vector(loops, [&](auto loop) {
- return clampSizeToNonNegative(rewriter, loc, loop.getUpperBound());
+ return clampSize(rewriter, loc, loop.getUpperBound(), /*minValue=*/0);
+ });
+ SmallVector<Value> basis = llvm::map_to_vector(loops, [&](auto loop) {
+ return clampSize(rewriter, loc, loop.getUpperBound(), /*minValue=*/1);
});
Value upperBound = getProductOfIntsOrIndexes(rewriter, loc, upperBounds);
outermost.setUpperBound(upperBound);
@@ -1059,7 +1060,7 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
// Insert delinearization at the start of the outermost loop body.
rewriter.setInsertionPointToStart(outermost.getBody());
auto [delinearizeIvs, preservedUsers] = delinearizeInductionVariable(
- rewriter, loc, outermost.getInductionVar(), upperBounds);
+ rewriter, loc, outermost.getInductionVar(), basis);
rewriter.replaceAllUsesExcept(outermost.getInductionVar(), delinearizeIvs[0],
preservedUsers);
@@ -1197,6 +1198,7 @@ void mlir::collapseParallelLoops(
// Normalize ParallelOp's iteration pattern.
SmallVector<Value, 3> normalizedUpperBounds;
+ SmallVector<Value, 3> normalizedDivisors;
for (unsigned i = 0, e = loops.getNumLoops(); i < e; ++i) {
OpBuilder::InsertionGuard g2(rewriter);
rewriter.setInsertionPoint(loops);
@@ -1204,10 +1206,12 @@ void mlir::collapseParallelLoops(
Value ub = loops.getUpperBound()[i];
Value step = loops.getStep()[i];
auto newLoopRange = emitNormalizedLoopBounds(rewriter, loc, lb, ub, step);
- normalizedUpperBounds.push_back(clampSizeToNonNegative(
- rewriter, loops.getLoc(),
- getValueOrCreateConstantIntOp(rewriter, loops.getLoc(),
- newLoopRange.size)));
+ Value size = getValueOrCreateConstantIntOp(rewriter, loops.getLoc(),
+ newLoopRange.size);
+ normalizedUpperBounds.push_back(
+ clampSize(rewriter, loops.getLoc(), size, /*minValue=*/0));
+ normalizedDivisors.push_back(
+ clampSize(rewriter, loops.getLoc(), size, /*minValue=*/1));
rewriter.setInsertionPointToStart(loops.getBody());
denormalizeInductionVariable(rewriter, loc, loops.getInductionVars()[i], lb,
@@ -1247,14 +1251,14 @@ void mlir::collapseParallelLoops(
// Determine the current induction value's current loop iteration
Value iv = arith::RemSIOp::create(insideBuilder, loc, previous,
- normalizedUpperBounds[idx]);
+ normalizedDivisors[idx]);
replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx), iv,
loops.getRegion());
// Remove the effect of the current induction value to prepare for
// the next value.
previous = arith::DivSIOp::create(insideBuilder, loc, previous,
- normalizedUpperBounds[idx]);
+ normalizedDivisors[idx]);
}
// The final induction value is just the remaining value.
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index b33a3b1efc01c..72de817cfd784 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -227,9 +227,12 @@ func.func @parametric(%lb1 : index, %ub1 : index, %step1 : index,
// CHECK: %[[normalized_j:.*]] = affine.apply
// CHECK-SAME: affine_map<()[s0, s1, s2] -> ((-s0 + s1) ceildiv s2)>()[%[[orig_lb2]], %[[orig_ub2]], %[[orig_step2]]]
// Each normalized size is clamped at zero before the sizes are combined, so
- // two empty dimensions cannot cancel into a positive flattened extent.
- // CHECK: %[[clamped_i:.+]] = arith.maxsi %[[normalized_i]]
- // CHECK: %[[clamped_j:.+]] = arith.maxsi %[[normalized_j]]
+ // two empty dimensions cannot cancel into a positive flattened extent, and at
+ // one where it is used as the delinearization basis.
+ // CHECK: %[[clamped_i:.+]] = arith.maxsi %[[normalized_i]], %[[c0]]
+ // CHECK: %[[clamped_j:.+]] = arith.maxsi %[[normalized_j]], %[[c0]]
+ // CHECK: %[[basis_i:.+]] = arith.maxsi %[[normalized_i]], %[[c1]]
+ // CHECK: %[[basis_j:.+]] = arith.maxsi %[[normalized_j]], %[[c1]]
// CHECK: %[[range:.+]] = affine.apply
// CHECK-SAME: affine_map<()[s0, s1] -> (s0 * s1)>()[%[[clamped_i]], %[[clamped_j]]]
@@ -240,7 +243,7 @@ func.func @parametric(%lb1 : index, %ub1 : index, %step1 : index,
// CHECK-NOT: scf.for
scf.for %j = %lb2 to %ub2 step %step2 {
// Remapping of the induction variables.
- // CHECK: %[[delinearize:.+]]:2 = affine.delinearize_index %[[i]] into (%[[clamped_i]], %[[clamped_j]])
+ // CHECK: %[[delinearize:.+]]:2 = affine.delinearize_index %[[i]] into (%[[basis_i]], %[[basis_j]])
// CHECK: %[[orig_j:.*]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
// CHECK-SAME: (%[[delinearize]]#1)[%[[orig_lb2]], %[[orig_step2]]]
// CHECK: %[[orig_i:.*]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
diff --git a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
index bbf7ff986bba8..69ba94d42a8a1 100644
--- a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
+++ b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
@@ -142,13 +142,17 @@ module attributes {transform.with_named_sequence} {
// CHECK: %[[NITERS2:.+]] = affine.apply
// CHECK-SAME: affine_map<()[s0, s1, s2] -> ((-s0 + s1) ceildiv s2)>()[%[[LB2]], %[[UB2]], %[[STEP2]]]
// Each size is clamped at zero before the sizes are combined, so empty
-// dimensions cannot cancel into a positive flattened extent.
-// CHECK: %[[CL0:.+]] = arith.maxsi %[[NITERS0]]
-// CHECK: %[[CL1:.+]] = arith.maxsi %[[NITERS1]]
-// CHECK: %[[CL2:.+]] = arith.maxsi %[[NITERS2]]
+// dimensions cannot cancel into a positive flattened extent, and at one where
+// it is used as the delinearization basis.
+// CHECK: %[[CL0:.+]] = arith.maxsi %[[NITERS0]], %[[C0]]
+// CHECK: %[[CL1:.+]] = arith.maxsi %[[NITERS1]], %[[C0]]
+// CHECK: %[[CL2:.+]] = arith.maxsi %[[NITERS2]], %[[C0]]
+// CHECK: %[[B0:.+]] = arith.maxsi %[[NITERS0]], %[[C1]]
+// CHECK: %[[B1:.+]] = arith.maxsi %[[NITERS1]], %[[C1]]
+// CHECK: %[[B2:.+]] = arith.maxsi %[[NITERS2]], %[[C1]]
// CHECK: %[[NEWUB:.+]] = affine.apply
// CHECK: %[[RESULT:.+]] = scf.for %[[IV:[a-zA-Z0-9]+]] = %[[C0]] to %[[NEWUB]] step %[[C1]] iter_args(%[[ITER_ARG:.+]] = %[[ARG0]])
-// CHECK: %[[DELINEARIZE:.+]]:3 = affine.delinearize_index %[[IV]] into (%[[CL0]], %[[CL1]], %[[CL2]])
+// CHECK: %[[DELINEARIZE:.+]]:3 = affine.delinearize_index %[[IV]] into (%[[B0]], %[[B1]], %[[B2]])
// CHECK-DAG: %[[K:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#2)[%[[LB2]], %[[STEP2]]]
// CHECK-DAG: %[[J:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#1)[%[[LB1]], %[[STEP1]]]
// CHECK-DAG: %[[I:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#0)[%[[LB0]], %[[STEP0]]]
@@ -327,11 +331,13 @@ module attributes {transform.with_named_sequence} {
// CHECK-SAME: %[[ARG2:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK: %[[CL1:.+]] = arith.maxsi %[[ARG1]]
-// CHECK: %[[CL2:.+]] = arith.maxsi %[[ARG2]]
+// CHECK: %[[CL1:.+]] = arith.maxsi %[[ARG1]], %[[C0]]
+// CHECK: %[[CL2:.+]] = arith.maxsi %[[ARG2]], %[[C0]]
+// CHECK: %[[B1:.+]] = arith.maxsi %[[ARG1]], %[[C1]]
+// CHECK: %[[B2:.+]] = arith.maxsi %[[ARG2]], %[[C1]]
// CHECK: %[[UB:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%[[CL1]], %[[CL2]]]
// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[UB]] step %[[C1]]
-// CHECK: %[[DELINEARIZE:.+]]:2 = affine.delinearize_index %[[IV]] into (%[[CL1]], %[[CL2]])
+// CHECK: %[[DELINEARIZE:.+]]:2 = affine.delinearize_index %[[IV]] into (%[[B1]], %[[B2]])
// CHECK: "some_use"(%{{[a-zA-Z0-9]+}}, %[[C0]], %[[C0]], %[[DELINEARIZE]]#0, %[[C0]], %[[DELINEARIZE]]#1)
// -----
@@ -368,7 +374,7 @@ module attributes {transform.with_named_sequence} {
// CHECK-SAME: , %[[ARG1:.+]]: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK: %[[CL:.+]] = arith.maxsi %[[ARG1]]
+// CHECK: %[[CL:.+]] = arith.maxsi %[[ARG1]], %[[C0]]
// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[CL]] step %[[C1]]
// CHECK: "some_use"(%{{[a-zA-Z0-9]+}}, %[[C0]], %[[C0]], %[[IV]])
@@ -420,3 +426,33 @@ module attributes {transform.with_named_sequence} {
// CHECK: scf.yield %[[UPDATED]]
// CHECK: }
// CHECK: return %[[RESULT]]
+
+// -----
+
+// Statically empty loops: the sizes are negative constants. The flattened
+// bound must fold to zero and the delinearization basis must stay legal.
+func.func @coalesce_empty_static_loops() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ scf.for %i = %c2 to %c0 step %c1 {
+ scf.for %j = %c2 to %c0 step %c1 {
+ "some_use"(%i, %j) : (index, index) -> ()
+ }
+ } {coalesce}
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
+ %2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
+ transform.yield
+ }
+}
+// CHECK-LABEL: func @coalesce_empty_static_loops
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[C0]] step %[[C1]]
+// CHECK: affine.delinearize_index %[[IV]] into (1, 1)
diff --git a/mlir/test/Transforms/single-parallel-loop-collapsing.mlir b/mlir/test/Transforms/single-parallel-loop-collapsing.mlir
index 1ef787bec1bb3..77d5452aecaa2 100644
--- a/mlir/test/Transforms/single-parallel-loop-collapsing.mlir
+++ b/mlir/test/Transforms/single-parallel-loop-collapsing.mlir
@@ -27,3 +27,26 @@ func.func @collapse_to_single() {
// CHECK: scf.reduce
// CHECK-NEXT: }
// CHECK-NEXT: return
+
+// Dynamic bounds: each normalized size is clamped at zero for the flattened
+// bound and at one where it is used as a divisor.
+// CHECK-LABEL: func @collapse_dynamic_bounds
+// CHECK-SAME: %[[LB0:.+]]: index, %[[UB0:.+]]: index, %[[LB1:.+]]: index, %[[UB1:.+]]: index)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK: %[[SIZE0:.+]] = affine.apply affine_map<()[s0, s1] -> (-s0 + s1)>()[%[[LB0]], %[[UB0]]]
+// CHECK: %[[CL0:.+]] = arith.maxsi %[[SIZE0]], %[[C0]]
+// CHECK: %[[SIZE1:.+]] = affine.apply affine_map<()[s0, s1] -> (-s0 + s1)>()[%[[LB1]], %[[UB1]]]
+// CHECK: %[[CL1:.+]] = arith.maxsi %[[SIZE1]], %[[C0]]
+// CHECK: %[[DIV1:.+]] = arith.maxsi %[[SIZE1]], %[[C1]]
+// CHECK: %[[UB:.+]] = arith.muli %[[CL0]], %[[CL1]]
+// CHECK: scf.parallel (%[[IV:.+]]) = (%[[C0]]) to (%[[UB]]) step (%[[C1]])
+// CHECK: arith.remsi %[[IV]], %[[DIV1]]
+// CHECK: arith.divsi %[[IV]], %[[DIV1]]
+func.func @collapse_dynamic_bounds(%lb0: index, %ub0: index, %lb1: index, %ub1: index) {
+ %c1 = arith.constant 1 : index
+ scf.parallel (%i0, %i1) = (%lb0, %lb1) to (%ub0, %ub1) step (%c1, %c1) {
+ %result = "magic.op"(%i0, %i1) : (index, index) -> index
+ }
+ return
+}
>From b79aba5f8eabaa5d7550f01799e1d3b07331f3eb Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Tue, 6 Oct 2026 00:52:31 +0000
Subject: [PATCH 3/3] Address review: reword the empty-loop comment and test a
non-zero lower bound
The clamp applies to the normalized size ceildiv(ub - lb, step), not to the
upper bound; the new test uses lb = -1 so the two differ.
---
.../Dialect/SCF/transform-op-coalesce.mlir | 43 ++++++++++++++++++-
1 file changed, 41 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
index 69ba94d42a8a1..49a5d52cf2e1c 100644
--- a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
+++ b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
@@ -429,8 +429,9 @@ module attributes {transform.with_named_sequence} {
// -----
-// Statically empty loops: the sizes are negative constants. The flattened
-// bound must fold to zero and the delinearization basis must stay legal.
+// Statically empty loops (ub < lb): the normalized sizes are negative
+// constants. The flattened bound must fold to zero and the delinearization
+// basis must stay legal.
func.func @coalesce_empty_static_loops() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
@@ -456,3 +457,41 @@ module attributes {transform.with_named_sequence} {
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[C0]] step %[[C1]]
// CHECK: affine.delinearize_index %[[IV]] into (1, 1)
+
+// -----
+
+// The clamp applies to the normalized size, not to the upper bound: with a
+// non-zero lower bound the size is ub - lb.
+func.func @coalesce_dynamic_bound_nonzero_lb(%arg0 : index) {
+ %c-1 = arith.constant -1 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ scf.for %i = %c-1 to %arg0 step %c1 {
+ scf.for %j = %c0 to %c4 step %c1 {
+ "some_use"(%i, %j) : (index, index) -> ()
+ }
+ } {coalesce}
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
+ %2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
+ transform.yield
+ }
+}
+// CHECK-LABEL: func @coalesce_dynamic_bound_nonzero_lb
+// CHECK-SAME: %[[ARG0:.+]]: index)
+// CHECK-DAG: %[[SIZE:.+]] = affine.apply affine_map<()[s0] -> (s0 + 1)>()[%[[ARG0]]]
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK: %[[CL:.+]] = arith.maxsi %[[SIZE]], %[[C0]]
+// CHECK: %[[B:.+]] = arith.maxsi %[[SIZE]], %[[C1]]
+// CHECK: %[[UB:.+]] = affine.apply affine_map<()[s0] -> (s0 * 4)>()[%[[CL]]]
+// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[UB]] step %[[C1]]
+// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[IV]] into (%[[B]], 4)
+// CHECK: %[[I:.+]] = affine.apply affine_map<(d0) -> (d0 - 1)>(%[[DELIN]]#0)
+// CHECK: "some_use"(%[[I]], %[[DELIN]]#1)
More information about the Mlir-commits
mailing list