Skip to content

Commit

Permalink
[flang][OpenMP] Support host_eval for target teams loop
Browse files Browse the repository at this point in the history
Extends `host_eval` support for the currently supported form of the
generic `loop` directive.
  • Loading branch information
ergawy committed Jan 13, 2025
1 parent 4f083b3 commit 8e21c47
Showing 4 changed files with 44 additions and 9 deletions.
16 changes: 14 additions & 2 deletions flang/lib/Lower/OpenMP/OpenMP.cpp
Original file line number Diff line number Diff line change
@@ -567,8 +567,9 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
[[fallthrough]];
case OMPD_target_teams:
cp.processNumTeams(stmtCtx, hostInfo.ops);
processSingleNestedIf(
[](Directive nestedDir) { return topDistributeSet.test(nestedDir); });
processSingleNestedIf([](Directive nestedDir) {
return topDistributeSet.test(nestedDir) || topLoopSet.test(nestedDir);
});
break;

case OMPD_teams_distribute:
@@ -581,6 +582,17 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
cp.processNumTeams(stmtCtx, hostInfo.ops);
break;


case OMPD_teams_loop:
cp.processThreadLimit(stmtCtx, hostInfo.ops);
[[fallthrough]];
case OMPD_target_teams_loop:
cp.processNumTeams(stmtCtx, hostInfo.ops);
[[fallthrough]];
case OMPD_loop:
cp.processCollapse(loc, eval, hostInfo.ops, hostInfo.iv);
break;

// Standalone 'target' case.
case OMPD_target:
processSingleNestedIf(
12 changes: 7 additions & 5 deletions flang/test/Lower/OpenMP/generic-loop-rewriting.f90
Original file line number Diff line number Diff line change
@@ -11,7 +11,13 @@ subroutine target_teams_loop
end subroutine target_teams_loop

!CHECK-LABEL: func.func @_QPtarget_teams_loop
!CHECK: omp.target map_entries(
!CHECK: omp.target
!CHECK-SAME: host_eval(
!CHECK-SAME: %c0{{[^[:space:]]+}} -> %[[LB:[^[:space:]]+]],
!CHECK-SAME: %c10{{[^[:space:]]+}} -> %[[UB:[^[:space:]]+]],
!CHECK-SAME: %c1{{[^[:space:]]+}} -> %[[STEP:[^[:space:]]+]]
!CHECK-SAME: : i32, i32, i32)
!CHECK-SAME: map_entries(
!CHECK-SAME: %{{.*}} -> %[[I_ARG:[^[:space:]]+]],
!CHECK-SAME: %{{.*}} -> %[[X_ARG:[^[:space:]]+]] : {{.*}}) {

@@ -20,10 +26,6 @@ end subroutine target_teams_loop

!CHECK: omp.teams {

!CHECK: %[[LB:.*]] = arith.constant 0 : i32
!CHECK: %[[UB:.*]] = arith.constant 10 : i32
!CHECK: %[[STEP:.*]] = arith.constant 1 : i32

!CHECK: omp.parallel private(@{{.*}} %[[I_DECL]]#0
!CHECK-SAME: -> %[[I_PRIV_ARG:[^[:space:]]+]] : !fir.ref<i32>) {
!CHECK: omp.distribute {
7 changes: 5 additions & 2 deletions mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
Original file line number Diff line number Diff line change
@@ -1908,16 +1908,19 @@ llvm::omp::OMPTgtExecModeFlags TargetOp::getKernelExecFlags() {
long numWrappers = std::distance(innermostWrapper, wrappers.end());

// Detect Generic-SPMD: target-teams-distribute[-simd].
// Detect SPMD: target-teams-loop.
if (numWrappers == 1) {
if (!isa<DistributeOp>(innermostWrapper))
if (!isa<DistributeOp, LoopOp>(innermostWrapper))
return OMP_TGT_EXEC_MODE_GENERIC;

Operation *teamsOp = (*innermostWrapper)->getParentOp();
if (!isa_and_present<TeamsOp>(teamsOp))
return OMP_TGT_EXEC_MODE_GENERIC;

if (teamsOp->getParentOp() == *this)
return OMP_TGT_EXEC_MODE_GENERIC_SPMD;
return isa<DistributeOp>(innermostWrapper)
? OMP_TGT_EXEC_MODE_GENERIC_SPMD
: OMP_TGT_EXEC_MODE_SPMD;
}

// Detect SPMD: target-teams-distribute-parallel-wsloop[-simd].
18 changes: 18 additions & 0 deletions offload/test/offloading/fortran/target_teams_loop.f90
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
! REQUIRES: flang, amdgpu

! RUN: %libomptarget-compile-fortran-generic
! RUN: env LIBOMPTARGET_INFO=16 %libomptarget-run-generic 2>&1 | %fcheck-generic
program target_teams_loop
implicit none
integer :: x(10), i

!$omp target teams loop
do i = 1, 10
x(i) = i * 2
end do

print *, x
end program target_teams_loop

! CHECK: "PluginInterface" device {{[0-9]+}} info: Launching kernel {{.*}}
! CHECK: 2 4 6 8 10 12 14 16 18 20

0 comments on commit 8e21c47

Please sign in to comment.