[Mlir-commits] [mlir] [mlir][async] Clamp per-axis trip counts in async-parallel-for (PR #225370)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Nov 1 04:32:32 PST 2026
https://github.com/Shoryamishra61 updated https://github.com/llvm/llvm-project/pull/225370
>From 19a89ee55353105c1f203f1b1241e00b677f4663 Mon Sep 17 00:00:00 2001
From: Shorya Mishra <shoryamishra181 at gmail.com>
Date: Tue, 22 Sep 2026 16:49:02 +0530
Subject: [PATCH 1/3] [mlir][async] Clamp per-axis trip counts in
async-parallel-for
AsyncParallelFor derives per-dimension trip counts using signed arithmetic.
An empty dimension with a positive step produces a negative raw trip count.
For a multidimensional loop with multiple empty dimensions, multiplying
negative counts can cancel out signs and produce a positive flattened
iteration count, causing an empty scf.parallel domain to execute work.
Normalize each per-dimension trip count to a non-negative cardinality
using arith::MaxSIOp before computing the product.
Fixes #225254
---
.../Async/Transforms/AsyncParallelFor.cpp | 10 +++--
.../async-parallel-for-async-dispatch.mlir | 40 +++++++++++++++++++
2 files changed, 47 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
index 1db6f128f4c697..70c705b294d447 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
@@ -723,14 +723,19 @@ AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
cloneConstantsIntoTheRegion(op.getRegion(), rewriter);
// Compute trip count for each loop induction variable:
- // tripCount = ceil_div(upperBound - lowerBound, step);
+ // tripCount = max(ceil_div(upperBound - lowerBound, step), 0);
+ // Each dimension is clamped at zero so that empty dimensions (where
+ // upperBound < lowerBound) contribute zero iterations rather than negative
+ // values which can cancel when multiplied together.
+ Value c0 = arith::ConstantIndexOp::create(b, 0);
SmallVector<Value> tripCounts(op.getNumLoops());
for (size_t i = 0; i < op.getNumLoops(); ++i) {
auto lb = op.getLowerBound()[i];
auto ub = op.getUpperBound()[i];
auto step = op.getStep()[i];
auto range = b.createOrFold<arith::SubIOp>(ub, lb);
- tripCounts[i] = b.createOrFold<arith::CeilDivSIOp>(range, step);
+ Value rawTripCount = b.createOrFold<arith::CeilDivSIOp>(range, step);
+ tripCounts[i] = b.createOrFold<arith::MaxSIOp>(rawTripCount, c0);
}
// Compute a product of trip counts to get the 1-dimensional iteration space
@@ -741,7 +746,6 @@ AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
// Short circuit no-op parallel loops (zero iterations) that can arise from
// the memrefs with dynamic dimension(s) equal to zero.
- Value c0 = arith::ConstantIndexOp::create(b, 0);
Value isZeroIterations =
arith::CmpIOp::create(b, arith::CmpIPredicate::eq, tripCount, c0);
diff --git a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
index 3f412ecd0c8f9b..8359e072d6dfeb 100644
--- a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
+++ b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
@@ -84,3 +84,43 @@ func.func @loop_2d(%arg0: index, %arg1: index, %arg2: index, // lb, ub, step
// CHECK: memref.store
// CHECK-LABEL: func private @async_dispatch_fn
+
+// -----
+
+// CHECK-LABEL: @empty_2d_dynamic
+// CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
+func.func @empty_2d_dynamic(%lb: index, %ub: index, %step: index, %arg0: memref<?x?xf32>) {
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: %[[RANGE0:.*]] = arith.subi %[[UB]], %[[LB]]
+ // CHECK: %[[RAW0:.*]] = arith.ceildivsi %[[RANGE0]], %[[STEP]]
+ // CHECK: %[[TRIP0:.*]] = arith.maxsi %[[RAW0]], %[[C0]]
+ // CHECK: %[[RANGE1:.*]] = arith.subi %[[UB]], %[[LB]]
+ // CHECK: %[[RAW1:.*]] = arith.ceildivsi %[[RANGE1]], %[[STEP]]
+ // CHECK: %[[TRIP1:.*]] = arith.maxsi %[[RAW1]], %[[C0]]
+ // CHECK: %[[FLATTENED:.*]] = arith.muli %[[TRIP0]], %[[TRIP1]]
+ // CHECK: %[[IS_NOOP:.*]] = arith.cmpi eq, %[[FLATTENED]], %[[C0]] : index
+ // CHECK: scf.if %[[IS_NOOP]] {
+ // CHECK-NEXT: } else {
+ scf.parallel (%i, %j) = (%lb, %lb) to (%ub, %ub) step (%step, %step) {
+ %one = arith.constant 1.0 : f32
+ memref.store %one, %arg0[%i, %j] : memref<?x?xf32>
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @empty_2d_static
+func.func @empty_2d_static(%arg0: memref<8x8xi32>) {
+ // CHECK: %[[TRUE:.*]] = arith.constant true
+ // CHECK: scf.if %[[TRUE]] {
+ // CHECK-NEXT: } else {
+ %c5 = arith.constant 5 : index
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %one = arith.constant 1 : i32
+ scf.parallel (%i, %j) = (%c5,%c5) to (%c0,%c0) step (%c2,%c2) {
+ memref.store %one, %arg0[%i, %j] : memref<8x8xi32>
+ }
+ return
+}
>From 59f419a99cc7547a43e212b9e9e50a70a08b73c5 Mon Sep 17 00:00:00 2001
From: Shorya Mishra <shoryamishra181 at gmail.com>
Date: Tue, 22 Sep 2026 17:09:43 +0530
Subject: [PATCH 2/3] [mlir][async][test] Update FileCheck expectations for
clamped trip counts
- Update loop_1d to check that cmpi consumes the clamped maxsi trip count rather than the raw ceildivsi result.
- Update empty_2d_static to match the non-canonicalized muli/cmpi IR produced by the pass without requiring premature constant folding.
---
.../Dialect/Async/async-parallel-for-async-dispatch.mlir | 8 +++++---
1 file changed, 5 insertions(+), 3 deletions(-)
diff --git a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
index 8359e072d6dfeb..f40449e7f016a6 100644
--- a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
+++ b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
@@ -7,7 +7,8 @@ func.func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[RANGE:.*]] = arith.subi %[[UB]], %[[LB]]
- // CHECK: %[[TRIP_CNT:.*]] = arith.ceildivsi %[[RANGE]], %[[STEP]]
+ // CHECK: %[[RAW_TRIP_CNT:.*]] = arith.ceildivsi %[[RANGE]], %[[STEP]]
+ // CHECK: %[[TRIP_CNT:.*]] = arith.maxsi %[[RAW_TRIP_CNT]], %[[C0]]
// CHECK: %[[IS_NOOP:.*]] = arith.cmpi eq, %[[TRIP_CNT]], %[[C0]] : index
// CHECK: scf.if %[[IS_NOOP]] {
@@ -112,8 +113,9 @@ func.func @empty_2d_dynamic(%lb: index, %ub: index, %step: index, %arg0: memref<
// CHECK-LABEL: @empty_2d_static
func.func @empty_2d_static(%arg0: memref<8x8xi32>) {
- // CHECK: %[[TRUE:.*]] = arith.constant true
- // CHECK: scf.if %[[TRUE]] {
+ // CHECK: %[[PROD:.*]] = arith.muli %c0, %c0 : index
+ // CHECK: %[[IS_ZERO:.*]] = arith.cmpi eq, %[[PROD]], %c0 : index
+ // CHECK: scf.if %[[IS_ZERO]] {
// CHECK-NEXT: } else {
%c5 = arith.constant 5 : index
%c0 = arith.constant 0 : index
>From dc333339f9d7762111888e98866dd92c1b546fcf Mon Sep 17 00:00:00 2001
From: Shoryamishra61 <shoryamishra65 at gmail.com>
Date: Sun, 1 Nov 2026 05:31:40 +0530
Subject: [PATCH 3/3] Restore FileCheck expectations for @loop_1d clamped trip
count
---
mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
index 96f21323121d22..8fbf0434b00abb 100644
--- a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
+++ b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
@@ -7,7 +7,8 @@ func.func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[RANGE:.*]] = arith.subi %[[UB]], %[[LB]]
- // CHECK: %[[TRIP_CNT:.*]] = arith.ceildivsi %[[RANGE]], %[[STEP]]
+ // CHECK: %[[RAW_TRIP_CNT:.*]] = arith.ceildivsi %[[RANGE]], %[[STEP]]
+ // CHECK: %[[TRIP_CNT:.*]] = arith.maxsi %[[RAW_TRIP_CNT]], %[[C0]]
// CHECK: %[[IS_NOOP:.*]] = arith.cmpi eq, %[[TRIP_CNT]], %[[C0]] : index
// CHECK: scf.if %[[IS_NOOP]] {
More information about the Mlir-commits
mailing list