Skip to content
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

[MLIR][Flang][OpenMP] Host evaluation of trip count for Generic-SPMD #213

Merged
merged 1 commit into from
Nov 27, 2024
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
20 changes: 9 additions & 11 deletions flang/lib/Lower/OpenMP/OpenMP.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -578,7 +578,6 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
HostEvalInfo &hostInfo = hostEvalInfo.back();

switch (extractOmpDirective(*ompEval)) {
// Cases where 'teams' and target SPMD clauses might be present.
case OMPD_teams_distribute_parallel_do:
case OMPD_teams_distribute_parallel_do_simd:
cp.processThreadLimit(stmtCtx, hostInfo.ops);
Expand All @@ -589,39 +588,38 @@ 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
// possible by looking at nested evaluations.
case OMPD_teams:
cp.processThreadLimit(stmtCtx, hostInfo.ops);
[[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.
case OMPD_teams_distribute:
case OMPD_teams_distribute_simd:
cp.processThreadLimit(stmtCtx, hostInfo.ops);
[[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;

// Standalone 'target' case.
case OMPD_target: {
case OMPD_target:
processSingleNestedIf(
[](Directive nestedDir) { return topTeamsSet.test(nestedDir); });
break;
}

default:
break;
}
Expand Down
103 changes: 103 additions & 0 deletions flang/test/Lower/OpenMP/eval-outside-target.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
191 changes: 191 additions & 0 deletions flang/test/Lower/OpenMP/target-generic-spmd.f90
Original file line number Diff line number Diff line change
@@ -0,0 +1,191 @@
! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s

! CHECK-LABEL: func.func @_QPdistribute_generic() {
subroutine distribute_generic()
! CHECK: omp.target
! CHECK-NOT: host_eval({{.*}})
! CHECK-SAME: {
!$omp target
!$omp teams
!$omp distribute
do i = 1, 10
call foo(i)
end do
!$omp end distribute
call bar() !< Prevents this from being Generic-SPMD.
!$omp end teams
!$omp end target

! CHECK: omp.target
! CHECK-NOT: host_eval({{.*}})
! CHECK-SAME: {
!$omp target teams
!$omp distribute
do i = 1, 10
call foo(i)
end do
!$omp end distribute
call bar() !< Prevents this from being Generic-SPMD.
!$omp end target teams

! CHECK: omp.target
! CHECK-NOT: host_eval({{.*}})
! CHECK-SAME: {
!$omp target teams
!$omp distribute
do i = 1, 10
call foo(i)
end do
!$omp end distribute

!$omp distribute
do i = 1, 10
call foo(i)
end do
!$omp end distribute
!$omp end target teams
end subroutine distribute_generic

! CHECK-LABEL: func.func @_QPdistribute_spmd() {
subroutine distribute_spmd()
! CHECK: omp.target
! CHECK-SAME: host_eval({{.*}})
!$omp target
!$omp teams
!$omp distribute
do i = 1, 10
call foo(i)
end do
!$omp end distribute
!$omp end teams
!$omp end target

! CHECK: omp.target
! CHECK-SAME: host_eval({{.*}})
!$omp target teams
!$omp distribute
do i = 1, 10
call foo(i)
end do
!$omp end distribute
!$omp end target teams
end subroutine distribute_spmd

! CHECK-LABEL: func.func @_QPdistribute_simd_generic() {
subroutine distribute_simd_generic()
! CHECK: omp.target
! CHECK-NOT: host_eval({{.*}})
! CHECK-SAME: {
!$omp target
!$omp teams
!$omp distribute simd
do i = 1, 10
call foo(i)
end do
!$omp end distribute simd
call bar() !< Prevents this from being Generic-SPMD.
!$omp end teams
!$omp end target

! CHECK: omp.target
! CHECK-NOT: host_eval({{.*}})
! CHECK-SAME: {
!$omp target teams
!$omp distribute simd
do i = 1, 10
call foo(i)
end do
!$omp end distribute simd
call bar() !< Prevents this from being Generic-SPMD.
!$omp end target teams

! CHECK: omp.target
! CHECK-NOT: host_eval({{.*}})
! CHECK-SAME: {
!$omp target teams
!$omp distribute simd
do i = 1, 10
call foo(i)
end do
!$omp end distribute simd

!$omp distribute simd
do i = 1, 10
call foo(i)
end do
!$omp end distribute simd
!$omp end target teams
end subroutine distribute_simd_generic

! CHECK-LABEL: func.func @_QPdistribute_simd_spmd() {
subroutine distribute_simd_spmd()
! CHECK: omp.target
! CHECK-SAME: host_eval({{.*}})
!$omp target
!$omp teams
!$omp distribute simd
do i = 1, 10
call foo(i)
end do
!$omp end distribute simd
!$omp end teams
!$omp end target

! CHECK: omp.target
! CHECK-SAME: host_eval({{.*}})
!$omp target teams
!$omp distribute simd
do i = 1, 10
call foo(i)
end do
!$omp end distribute simd
!$omp end target teams
end subroutine distribute_simd_spmd

! CHECK-LABEL: func.func @_QPteams_distribute_spmd() {
subroutine teams_distribute_spmd()
! CHECK: omp.target
! CHECK-SAME: host_eval({{.*}})
!$omp target
!$omp teams distribute
do i = 1, 10
call foo(i)
end do
!$omp end teams distribute
!$omp end target
end subroutine teams_distribute_spmd

! CHECK-LABEL: func.func @_QPteams_distribute_simd_spmd() {
subroutine teams_distribute_simd_spmd()
! CHECK: omp.target
! CHECK-SAME: host_eval({{.*}})
!$omp target
!$omp teams distribute simd
do i = 1, 10
call foo(i)
end do
!$omp end teams distribute simd
!$omp end target
end subroutine teams_distribute_simd_spmd

! CHECK-LABEL: func.func @_QPtarget_teams_distribute_spmd() {
subroutine target_teams_distribute_spmd()
! CHECK: omp.target
! CHECK-SAME: host_eval({{.*}})
!$omp target teams distribute
do i = 1, 10
call foo(i)
end do
!$omp end target teams distribute
end subroutine target_teams_distribute_spmd

! CHECK-LABEL: func.func @_QPtarget_teams_distribute_simd_spmd() {
subroutine target_teams_distribute_simd_spmd()
! CHECK: omp.target
! CHECK-SAME: host_eval({{.*}})
!$omp target teams distribute simd
do i = 1, 10
call foo(i)
end do
!$omp end target teams distribute simd
end subroutine target_teams_distribute_simd_spmd
8 changes: 4 additions & 4 deletions flang/test/Lower/OpenMP/target-spmd.f90
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ subroutine distribute_parallel_do_generic()
call foo(i)
end do
!$omp end distribute parallel do
call bar()
call bar() !< Prevents this from being SPMD.
!$omp end teams
!$omp end target

Expand All @@ -25,7 +25,7 @@ subroutine distribute_parallel_do_generic()
call foo(i)
end do
!$omp end distribute parallel do
call bar()
call bar() !< Prevents this from being SPMD.
!$omp end target teams

! CHECK: omp.target
Expand Down Expand Up @@ -83,7 +83,7 @@ subroutine distribute_parallel_do_simd_generic()
call foo(i)
end do
!$omp end distribute parallel do simd
call bar()
call bar() !< Prevents this from being SPMD.
!$omp end teams
!$omp end target

Expand All @@ -96,7 +96,7 @@ subroutine distribute_parallel_do_simd_generic()
call foo(i)
end do
!$omp end distribute parallel do simd
call bar()
call bar() !< Prevents this from being SPMD.
!$omp end target teams

! CHECK: omp.target
Expand Down
Loading