Skip to content

[Flang][OpenMP] Allow host evaluation of loop bounds for distribute #127822

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 25, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions flang/lib/Lower/OpenMP/OpenMP.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -562,8 +562,11 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
[[fallthrough]];
case OMPD_distribute_parallel_do:
case OMPD_distribute_parallel_do_simd:
cp.processCollapse(loc, eval, hostInfo.ops, hostInfo.iv);
cp.processNumThreads(stmtCtx, hostInfo.ops);
[[fallthrough]];
case OMPD_distribute:
case OMPD_distribute_simd:
cp.processCollapse(loc, eval, hostInfo.ops, hostInfo.iv);
break;

// Cases where 'teams' clauses might be present, and target SPMD is
Expand All @@ -573,10 +576,8 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
[[fallthrough]];
case OMPD_target_teams:
cp.processNumTeams(stmtCtx, hostInfo.ops);
processSingleNestedIf([](Directive nestedDir) {
return nestedDir == OMPD_distribute_parallel_do ||
nestedDir == OMPD_distribute_parallel_do_simd;
});
processSingleNestedIf(
[](Directive nestedDir) { return topDistributeSet.test(nestedDir); });
break;

// Cases where only 'teams' host-evaluated clauses might be present.
Expand All @@ -586,6 +587,7 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
[[fallthrough]];
case OMPD_target_teams_distribute:
case OMPD_target_teams_distribute_simd:
cp.processCollapse(loc, eval, hostInfo.ops, hostInfo.iv);
cp.processNumTeams(stmtCtx, hostInfo.ops);
break;

Expand Down
103 changes: 103 additions & 0 deletions flang/test/Lower/OpenMP/host-eval.f90
Original file line number Diff line number Diff line change
Expand Up @@ -155,3 +155,106 @@ subroutine distribute_parallel_do_simd()
!$omp end distribute parallel do simd
!$omp end teams
end subroutine distribute_parallel_do_simd

! BOTH-LABEL: func.func @_QPdistribute
subroutine distribute()
! BOTH: omp.target

! HOST-SAME: host_eval(%{{.*}} -> %[[LB:.*]], %{{.*}} -> %[[UB:.*]], %{{.*}} -> %[[STEP:.*]] : i32, i32, i32)

! DEVICE-NOT: host_eval({{.*}})
! DEVICE-SAME: {

! BOTH: omp.teams
!$omp target teams

! BOTH: omp.distribute
! BOTH-NEXT: omp.loop_nest

! HOST-SAME: (%{{.*}}) : i32 = (%[[LB]]) to (%[[UB]]) inclusive step (%[[STEP]])
!$omp distribute
do i=1,10
call foo()
end do
!$omp end distribute
!$omp end target teams

! BOTH: omp.target
! BOTH-NOT: host_eval({{.*}})
! BOTH-SAME: {
! BOTH: omp.teams
!$omp target teams
call foo() !< Prevents this from being Generic-SPMD.

! BOTH: omp.distribute
!$omp distribute
do i=1,10
call foo()
end do
!$omp end distribute
!$omp end target teams

! BOTH: omp.teams
!$omp teams

! BOTH: omp.distribute
!$omp distribute
do i=1,10
call foo()
end do
!$omp end distribute
!$omp end teams
end subroutine distribute

! BOTH-LABEL: func.func @_QPdistribute_simd
subroutine distribute_simd()
! BOTH: omp.target

! HOST-SAME: host_eval(%{{.*}} -> %[[LB:.*]], %{{.*}} -> %[[UB:.*]], %{{.*}} -> %[[STEP:.*]] : i32, i32, i32)

! DEVICE-NOT: host_eval({{.*}})
! DEVICE-SAME: {

! BOTH: omp.teams
!$omp target teams

! BOTH: omp.distribute
! BOTH-NEXT: omp.simd
! BOTH-NEXT: omp.loop_nest

! HOST-SAME: (%{{.*}}) : i32 = (%[[LB]]) to (%[[UB]]) inclusive step (%[[STEP]])
!$omp distribute simd
do i=1,10
call foo()
end do
!$omp end distribute simd
!$omp end target teams

! BOTH: omp.target
! BOTH-NOT: host_eval({{.*}})
! BOTH-SAME: {
! BOTH: omp.teams
!$omp target teams
call foo() !< Prevents this from being Generic-SPMD.

! BOTH: omp.distribute
! BOTH-NEXT: omp.simd
!$omp distribute simd
do i=1,10
call foo()
end do
!$omp end distribute simd
!$omp end target teams

! BOTH: omp.teams
!$omp teams

! BOTH: omp.distribute
! BOTH-NEXT: omp.simd
!$omp distribute simd
do i=1,10
call foo()
end do
!$omp end distribute simd
!$omp end teams
end subroutine distribute_simd
Loading