[Mlir-commits] [mlir] [mlir][async] Clamp per-axis trip counts in async-parallel-for (PR #225370)

[llvmlistbot at llvm.org]

https://github.com/Shoryamishra61 created https://github.com/llvm/llvm-project/pull/225370

### Problem

`AsyncParallelFor` computes each dimension's trip count using signed ceiling division and directly multiplies those values to form the flattened iteration count.

For an empty dimension with a positive step, the raw result can be negative. If multiple dimensions are empty, multiplying an even number of negative trip counts can produce a positive flattened count, causing work to be dispatched for an `scf.parallel` operation whose semantic iteration space is empty.

### Fix

Normalize each per-dimension trip count to a non-negative iteration cardinality before computing the flattened trip count:
```cpp
Value rawTripCount = b.createOrFold<arith::CeilDivSIOp>(range, step);
tripCounts[i] = b.createOrFold<arith::MaxSIOp>(rawTripCount, c0);
```

This ensures any empty dimension contributes zero to the Cartesian-product iteration space while preserving positive trip counts for non-empty dimensions.

### Tests

Add regression coverage to `mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir`:
- `@empty_2d_dynamic`: Multidimensional empty iteration domain with dynamic bounds ensuring each dimension clamps to zero and takes the no-op branch.
- `@empty_2d_static`: Constant bounds verifying compile-time folding of empty dimensions to a no-op branch.

Fixes #225254.


>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] [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 1db6f128f4c69..70c705b294d44 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 3f412ecd0c8f9..8359e072d6dfe 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
+}