[Mlir-commits] [mlir] [mlir][scf] Clamp normalized loop sizes where they are combined (PR #223256)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Sep 13 10:21:32 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine
Author: Samarth Narang (snarang181)
<details>
<summary>Changes</summary>
`emitNormalizedLoopBounds` returns `ceilDiv(ub - lb, step)`, which is negative for an empty loop. A single negative size is harmless 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.
This PR clamps at the points where sizes are combined rather than inside `emitNormalizedLoopBounds`, so loops that are merely normalized keep their bound unchanged.
Fixes #<!-- -->223231.
---
Full diff: https://github.com/llvm/llvm-project/pull/223256.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+22-4)
- (modified) mlir/test/Dialect/Affine/loop-coalescing.mlir (+6-3)
- (modified) mlir/test/Dialect/SCF/transform-op-coalesce.mlir (+13-7)
``````````diff
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]])
// -----
``````````
</details>
https://github.com/llvm/llvm-project/pull/223256
More information about the Mlir-commits
mailing list