diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp index fc69c6b870adf3..c54770c7fbfaee 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 visitedOps; + // Walk the parent region in pre-order to make sure we visit `targetOp` // before its nested ops. targetParentRegion->walk( @@ -149,6 +151,10 @@ class HostClausesInsertionGuard { if (operandDefiningOp == nullptr) continue; + if (visitedOps.contains(operandDefiningOp)) + continue; + + visitedOps.insert(operandDefiningOp); auto parentTargetOp = operandDefiningOp->getParentOfType(); @@ -179,7 +185,9 @@ class HostClausesInsertionGuard { mlir::IRMapping mapper; builder.setInsertionPoint(escapingOperand->getOwner()); + mlir::Operation *lastSliceOp = nullptr; + llvm::SetVector 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(op)) + if (llvm::isa(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.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: }