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

[Samarth Narang]

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

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.

The PR adjusts 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.

>From 7449f73eb87eb86af0e3f00e2116978072f5cf03 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 4da7ca8868655..c4c74d9fb9b87 100644
--- a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
@@ -255,3 +255,24 @@ func.func @orphan_loop_unroll_annotation(%buf: memref<8xi32>) {
   } independent attributes {llvm.loop_annotation = #llvm.loop_annotation<unroll = <disable = false, full = true>>} 
   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);