[flang][OpenMP] Handle usage of array elements in loop-control expres…

…sions

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

@@ -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;
  }

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: }