[Mlir-commits] [mlir] [mlir][scf] Clamp normalized loop sizes where they are combined (PR #223256)
Samarth Narang
llvmlistbot at llvm.org
Thu Sep 24 05:05:40 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/2] [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 7ba0a796c86315..92a969b10091ba 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 b8a68b9c5c1f98..b33a3b1efc01c0 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 b487b62f0e3fe6..bbf7ff986bba84 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/2] 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 92a969b10091ba..1d8a9641d11fcd 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 b33a3b1efc01c0..72de817cfd7840 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 bbf7ff986bba84..69ba94d42a8a15 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 1ef787bec1bb37..77d5452aecaa2c 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
+}
More information about the Mlir-commits
mailing list