[Mlir-commits] [mlir] [mlir][openacc] Adjust an exclusive upper bound in the direction of the step (PR #227678)

[Samarth Narang]

https://github.com/snarang181 updated https://github.com/llvm/llvm-project/pull/227678

>From 9240320ae919cc53fd33df7e1366f8f88ff50597 Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Wed, 30 Sep 2026 11:57:23 +0000
Subject: [PATCH] [mlir][openacc] Adjust an exclusive upper bound in the
 direction of the step

acc::calculateTripCount turned an exclusive upper bound into an inclusive one
by subtracting 1 regardless of the sign of the step. For a descending loop the
last iterated value lies above the exclusive bound, so the adjustment has to
add 1 instead: `10 to 1 step -3` iterates 10, 7 and 4 but was lowered to a
loop with four trips.

Adjust the bound by one unit towards the lower bound based on the sign of the
step. The select folds away for constant steps. The overflow of `ub - lb` for
extreme bounds reported in the same issue is not addressed here.

Addresses the step-direction half of #225255.
---
 .../OpenACC/Utils/OpenACCUtilsLoop.cpp        | 11 ++++++++--
 .../OpenACC/acc-compute-lowering-loop.mlir    | 21 +++++++++++++++++++
 .../Dialect/OpenACC/OpenACCUtilsLoopTest.cpp  | 12 +++++++++++
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
index 721762c960465..e936515aee3c1 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
@@ -26,7 +26,8 @@
 using namespace mlir;
 
 /// Calculate trip count for a loop: (ub - lb + step) / step
-/// If inclusiveUpperbound is false, subtracts 1 from ub first.
+/// If inclusiveUpperbound is false, first moves ub one unit towards lb, i.e.
+/// subtracts 1 for a positive step and adds 1 for a negative step.
 Value acc::calculateTripCount(OpBuilder &b, Location loc, Value lb, Value ub,
                               Value step, bool inclusiveUpperbound) {
   Type type = b.getIndexType();
@@ -37,8 +38,14 @@ Value acc::calculateTripCount(OpBuilder &b, Location loc, Value lb, Value ub,
   step = getValueOrCreateCastToIndexLike(b, loc, type, step);
 
   if (!inclusiveUpperbound) {
+    Value zero = arith::ConstantIndexOp::create(b, loc, 0);
     Value one = arith::ConstantIndexOp::create(b, loc, 1);
-    ub = b.createOrFold<arith::SubIOp>(loc, ub, one,
+    Value negOne = arith::ConstantIndexOp::create(b, loc, -1);
+    Value isDescending = b.createOrFold<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::slt, step, zero);
+    Value adjustment =
+        b.createOrFold<arith::SelectOp>(loc, isDescending, one, negOne);
+    ub = b.createOrFold<arith::AddIOp>(loc, ub, adjustment,
                                        arith::IntegerOverflowFlags::nsw);
   }
 
diff --git a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
index 1a529e8f102e8..2bda1e0826010 100644
--- a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
@@ -310,3 +310,24 @@ func.func @parallel_loop_gang_static_star(%buf: memref<1xi32>) {
   acc.copyout accPtr(%dev : memref<1xi32>) to varPtr(%buf : memref<1xi32>)
   return
 }
+
+// -----
+
+// The exclusive upper bound of a descending loop lies below the last
+// iteration: 10, 7 and 4 are iterated, so the normalized loop has three trips.
+// CHECK-LABEL: func.func @descending_loop_exclusive_bound
+func.func @descending_loop_exclusive_bound(%buf: memref<16xi32>) {
+  %lb = arith.constant 10 : index
+  %ub = arith.constant 1 : index
+  %step = arith.constant -3 : index
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK: %[[C3:.*]] = arith.constant 3 : index
+  // CHECK: scf.for %{{.*}} = %[[C0]] to %[[C3]] step %[[C1]]
+  acc.loop control(%i : index) = (%lb : index) to (%ub : index) step (%step : index) {
+    %vi = arith.index_cast %i : index to i32
+    memref.store %vi, %buf[%i] : memref<16xi32>
+    acc.yield
+  } independent
+  return
+}
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
index e4d94855a06ec..e462cd1a45854 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
@@ -1092,6 +1092,18 @@ TEST_F(OpenACCUtilsLoopTest, CalculateTripCountNegativeStep) {
   EXPECT_EQ(getConstantIndex(tripCount), 10);
 }
 
+TEST_F(OpenACCUtilsLoopTest,
+       CalculateTripCountNegativeStepExclusiveUpperBound) {
+  auto [module, funcOp] = createModuleWithFunc();
+
+  // 10, 7 and 4 are iterated; the exclusive bound 1 is not.
+  Value tripCount = acc::calculateTripCount(
+      b, loc, createIndexConstant(10), createIndexConstant(1),
+      createIndexConstant(-3), /*inclusiveUpperbound=*/false);
+
+  EXPECT_EQ(getConstantIndex(tripCount), 3);
+}
+
 TEST_F(OpenACCUtilsLoopTest, CalculateTripCountCastsOperandsToIndex) {
   SmallVector<Type> argTypes(3, b.getI32Type());
   auto [module, funcOp] = createModuleWithFuncArgs(argTypes);