From ad1cf4cf660516dc14a3af255ed07a9cb139cc97 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy@amd.com>
Date: Wed, 31 Jul 2024 23:04:21 -0500
Subject: [PATCH] [flang][OpenMP] Handle usage of array elements in
 loop-control expressions

Extends the fix-up logic for `trip_count` calculation in `target`
regions. Previously, if an array element was used to compute any of the
loop bounds, the trip-count calculation ops would extract arra elements
from the mapped declaration of the array inside the target region. This
commit hanles that situation.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp             | 17 ++++++++--
 .../OpenMP/target-do-loop-control-exprs.f90   | 31 +++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 9cefc54fee4597..4ee73645e5b967 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -134,6 +134,8 @@ class HostClausesInsertionGuard {
     assert(targetParentRegion != nullptr &&
            "Expected omp.target op to be nested in a parent region.");
 
+    llvm::DenseSet<mlir::Operation *> visitedOps;
+
     // Walk the parent region in pre-order to make sure we visit `targetOp`
     // before its nested ops.
     targetParentRegion->walk<mlir::WalkOrder::PreOrder>(
@@ -149,6 +151,10 @@ class HostClausesInsertionGuard {
             if (operandDefiningOp == nullptr)
               continue;
 
+            if (visitedOps.contains(operandDefiningOp))
+              continue;
+
+            visitedOps.insert(operandDefiningOp);
             auto parentTargetOp =
                 operandDefiningOp->getParentOfType<mlir::omp::TargetOp>();
 
@@ -179,7 +185,9 @@ class HostClausesInsertionGuard {
 
     mlir::IRMapping mapper;
     builder.setInsertionPoint(escapingOperand->getOwner());
+
     mlir::Operation *lastSliceOp = nullptr;
+    llvm::SetVector<mlir::Operation *> opsToClone;
 
     for (auto *op : backwardSlice) {
       // DeclareOps need special handling by searching for the corresponding ops
@@ -188,12 +196,17 @@ class HostClausesInsertionGuard {
       //
       // TODO this might need a more elaborate handling in the future but for
       // now this seems sufficient for our purposes.
-      if (llvm::isa<hlfir::DeclareOp>(op))
+      if (llvm::isa<hlfir::DeclareOp>(op)) {
+        opsToClone.clear();
         break;
+      }
 
-      lastSliceOp = builder.clone(*op, mapper);
+      opsToClone.insert(op);
     }
 
+    for (mlir::Operation *op : opsToClone)
+      lastSliceOp = builder.clone(*op, mapper);
+
     builder.restoreInsertionPoint(ip);
     return lastSliceOp;
   }
diff --git a/flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90 b/flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90
index 8f4813bb27cd12..b4d5cffffac1d6 100644
--- a/flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90
+++ b/flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90
@@ -62,3 +62,34 @@ subroutine foo_with_dummy_arg(nodes)
 ! CHECK:   omp.target
 
 ! CHECK: }
+
+
+subroutine bounds_expr_in_loop_control(array)
+  real, intent(out) :: array(:,:)
+  integer :: bounds(2), i, j
+  bounds = shape(array)
+
+  !$omp target teams distribute parallel do simd collapse(2)
+  do j = 1,bounds(2)
+    do i = 1,bounds(1)
+      array(i,j) = 0.
+    enddo
+  enddo
+end subroutine bounds_expr_in_loop_control
+
+
+! CHECK: func.func @_QPbounds_expr_in_loop_control(%[[FUNC_ARG:.*]]: {{.*}}) {
+
+! CHECK:   %[[BOUNDS_DECL:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "{{.*}}Ebounds"} : (!fir.ref<!fir.array<2xi32>>, !fir.shape<1>) -> ({{.*}})
+
+! Verify that the host declaration of `bounds` (i.e. not the target/mapped one)
+! is used for the trip count calculation. Trip count is calculation ops are emitted
+! directly before the `omp.target` op and after all `omp.map.info` op; hence the
+! `CHECK-NOT: ...` line.
+
+! CHECK: hlfir.designate %[[BOUNDS_DECL:.*]]#0 (%c2{{.*}})
+! CHECK: hlfir.designate %[[BOUNDS_DECL:.*]]#0 (%c1{{.*}})
+! CHECK-NOT: omp.map.info
+! CHECK: omp.target
+
+! CHECK: }