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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-openacc

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/227678.diff


3 Files Affected:

- (modified) mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp (+9-2) 
- (modified) mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir (+21) 
- (modified) mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp (+12) 


``````````diff
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);

``````````

</details>


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