Skip to content

Commit 82fe869

Browse files
edoyangoedoyango
committed
enable group update omp offload tests
Update test_mpp_domains to pass the OpenMP offload flag through the group update path and stage test data on the device for offloaded runs. Add explicit OpenMP maps and shared clauses for the offloaded setup and reference-copy loops, including shifted CGRID and BGRID extents. Unskip the group update OpenMP offload test in test_mpp_domains2.sh.
1 parent 418d557 commit 82fe869

2 files changed

Lines changed: 118 additions & 29 deletions

File tree

‎test_fms/mpp/test_mpp_domains.F90‎

Lines changed: 117 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ program test_mpp_domains
115115
test_group, test_global_sum, test_subset, test_nonsym_edge, &
116116
test_halosize_performance, test_adjoint, wide_halo, &
117117
test_unstruct, test_group_offload
118-
integer :: i, j, k, n, p
118+
integer :: i, j, k, n
119119
integer :: layout(2)
120120
integer :: id
121121
integer :: outunit, errunit, io_status
@@ -227,13 +227,8 @@ program test_mpp_domains
227227
endif
228228
if( test_group .or. test_group_offload) then
229229
if (mpp_pe() == mpp_root_pe()) print *, '--------------------> Testing group_update <-------------------'
230-
do p=1,factorial(3)
231-
call test_group_update_r4('Folded-north', p, omp_offload=test_group_offload)
232-
call test_group_update_r4('Cubic-Grid', p, omp_offload=test_group_offload)
233-
234-
call test_group_update_r8('Folded-north', p, omp_offload=test_group_offload)
235-
call test_group_update_r8('Cubic-Grid', p, omp_offload=test_group_offload)
236-
enddo
230+
call test_group_update( 'Folded-north', use_omp_offload=test_group_offload )
231+
call test_group_update( 'Cubic-Grid', use_omp_offload=test_group_offload )
237232
if (mpp_pe() == mpp_root_pe()) print *, '--------------------> Finished testing group_update <-------------------'
238233
endif
239234

@@ -2476,8 +2471,9 @@ subroutine test_mpp_global_sum( type )
24762471
end subroutine test_mpp_global_sum
24772472

24782473
!###############################################################
2479-
subroutine test_group_update( type )
2474+
subroutine test_group_update( type, use_omp_offload )
24802475
character(len=*), intent(in) :: type
2476+
logical, intent(in), optional :: use_omp_offload
24812477

24822478
type(domain2D) :: domain
24832479
integer :: num_contact, ntiles, npes_per_tile
@@ -2499,9 +2495,12 @@ subroutine test_group_update( type )
24992495
integer :: nx_save, ny_save
25002496
type(mpp_group_update_type) :: group_update
25012497
type(mpp_group_update_type), allocatable :: update_list(:)
2498+
logical :: do_omp_offload
25022499

25032500
folded_north = .false.
25042501
cubic_grid = .false.
2502+
do_omp_offload = .false.
2503+
if (present(use_omp_offload)) do_omp_offload = use_omp_offload
25052504

25062505
nx_save = nx
25072506
ny_save = ny
@@ -2609,6 +2608,8 @@ subroutine test_group_update( type )
26092608
allocate( y2(ism:iem, jsm:jem+shift, nz, num_fields) )
26102609
allocate( base(isc:iec+shift,jsc:jec+shift,nz) )
26112610
a1 = 0; x1 = 0; y1 = 0
2611+
! allocate a1, x1, y1 on GPU (GPU x1, y1 only used for CGRID test)
2612+
!$omp target enter data if(do_omp_offload) map(to: a1, x1, y1)
26122613

26132614
base = 0
26142615
do k = 1,nz
@@ -2624,9 +2625,16 @@ subroutine test_group_update( type )
26242625
call mpp_create_group_update(group_update, a1(:,:,:,l), domain, flags=WUPDATE+SUPDATE)
26252626
end do
26262627

2628+
! initialize a1 values on GPU only to check GPU data is being transferred
2629+
!$omp target teams loop if(do_omp_offload) default(shared) &
2630+
!$omp shared(num_fields, nz, jsc, jec, isc, iec, a1, base) map(tofrom: a1) map(to: base)
26272631
do l = 1, num_fields
2628-
a1(isc:iec,jsc:jec,:,l) = base(isc:iec,jsc:jec,:) + l*1e3
26292632
do k=1,nz
2633+
do j=jsc,jec
2634+
do i=isc,iec
2635+
a1(i,j,k,l) = base(i,j,k) + l*1e3
2636+
enddo
2637+
enddo
26302638
do i=isc-1,iec+1
26312639
a1(i,jsc-1,k,l) = 999;
26322640
a1(i,jec+1,k,l) = 999;
@@ -2637,14 +2645,19 @@ subroutine test_group_update( type )
26372645
enddo
26382646
enddo
26392647
enddo
2640-
2641-
a2 = a1
2642-
call mpp_do_group_update(group_update, domain, a1(isc,jsc,1,1))
2648+
!$omp target teams loop collapse(4) if(do_omp_offload) default(shared) &
2649+
!$omp shared(num_fields, nz, jsm, jem, ism, iem, a1, a2) map(to: a1) map(from: a2)
2650+
do l=1,num_fields ; do k=1,nz ; do j=jsm,jem ; do i=ism,iem
2651+
a2(i,j,k,l) = a1(i,j,k,l)
2652+
enddo ; enddo ; enddo ; enddo
2653+
call mpp_do_group_update(group_update, domain, a1(isc,jsc,1,1), omp_offload=do_omp_offload)
26432654

26442655
do l = 1, num_fields
26452656
call mpp_update_domains( a2(:,:,:,l), domain, flags=WUPDATE+SUPDATE, complete=l==num_fields )
26462657
enddo
26472658

2659+
! copy a1 back to host
2660+
!$omp target update if(do_omp_offload) from(a1)
26482661
do l = 1, num_fields
26492662
write(text, '(i3.3)') l
26502663
call compare_checksums(a1(isd:ied,jsd:jed,:,l),a2(isd:ied,jsd:jed,:,l),type//' CENTER South West '//text)
@@ -2689,14 +2702,54 @@ subroutine test_group_update( type )
26892702

26902703
do n = 1, num_iter
26912704
a1 = 0; x1 = 0; y1 = 0
2705+
!$omp target update if(do_omp_offload) to(a1, x1, y1)
26922706
do l = 1, num_fields
2693-
a1(isc:iec, jsc:jec, :,l) = base(isc:iec, jsc:jec, :) + l*1e3
2694-
x1(isc:iec+shift,jsc:jec, :,l) = base(isc:iec+shift,jsc:jec, :) + l*1e3 + 1e6
2695-
y1(isc:iec, jsc:jec+shift,:,l) = base(isc:iec, jsc:jec+shift,:) + l*1e3 + 2*1e6
2707+
!$omp target teams loop collapse(3) if(do_omp_offload) default(shared) &
2708+
!$omp shared(nz, jsc, jec, isc, iec, a1, base, l) map(to: base) map(tofrom: a1)
2709+
do k=1,nz
2710+
do j=jsc,jec
2711+
do i=isc,iec
2712+
a1(i,j,k,l) = base(i,j,k) + l*1e3
2713+
enddo
2714+
enddo
2715+
enddo
2716+
!$omp target teams loop collapse(3) if(do_omp_offload) default(shared) &
2717+
!$omp shared(nz, jsc, jec, shift, isc, iec, x1, base, l) map(to: base) map(tofrom: x1)
2718+
do k=1,nz
2719+
do j=jsc,jec
2720+
do i=isc,iec+shift
2721+
x1(i,j,k,l) = base(i,j,k) + l*1e3 + 1e6
2722+
enddo
2723+
enddo
2724+
enddo
2725+
!$omp target teams loop collapse(3) if(do_omp_offload) default(shared) &
2726+
!$omp shared(nz, jsc, jec, shift, isc, iec, y1, base, l) map(to: base) map(tofrom: y1)
2727+
do k=1,nz
2728+
do j=jsc,jec+shift
2729+
do i=isc,iec
2730+
y1(i,j,k,l) = base(i,j,k) + l*1e3 + 2*1e6
2731+
enddo
2732+
enddo
2733+
enddo
26962734
enddo
2697-
a2 = a1; x2 = x1; y2 = y1
2735+
!$omp target teams loop collapse(4) if(do_omp_offload) default(shared) &
2736+
!$omp shared(num_fields, nz, jsm, jem, ism, iem, a1, a2) map(from: a2) map(to: a1)
2737+
do l=1,num_fields ; do k=1,nz ; do j=jsm,jem ; do i=ism,iem
2738+
a2(i,j,k,l) = a1(i,j,k,l)
2739+
enddo ; enddo ; enddo ; enddo
2740+
!$omp target teams loop collapse(4) if(do_omp_offload) default(shared) &
2741+
!$omp shared(num_fields, nz, jsm, jem, ism, iem, shift, x1, x2) map(from: x2) map(to: x1)
2742+
do l=1,num_fields ; do k=1,nz ; do j=jsm,jem ; do i=ism,iem+shift
2743+
x2(i,j,k,l) = x1(i,j,k,l)
2744+
enddo ; enddo ; enddo ; enddo
2745+
!$omp target teams loop collapse(4) if(do_omp_offload) default(shared) &
2746+
!$omp shared(num_fields, nz, jsm, jem, shift, ism, iem, y1, y2) map(from: y2) map(to: y1)
2747+
do l=1,num_fields ; do k=1,nz ; do j=jsm,jem+shift ; do i=ism,iem
2748+
y2(i,j,k,l) = y1(i,j,k,l)
2749+
enddo ; enddo ; enddo ; enddo
2750+
26982751
call mpp_clock_begin(id1)
2699-
call mpp_do_group_update(group_update, domain, a1(isc,jsc,1,1))
2752+
call mpp_do_group_update(group_update, domain, a1(isc,jsc,1,1), omp_offload=do_omp_offload)
27002753
call mpp_clock_end (id1)
27012754

27022755
call mpp_clock_begin(id2)
@@ -2710,6 +2763,7 @@ subroutine test_group_update( type )
27102763

27112764
!--- compare checksum
27122765
if( n == num_iter ) then
2766+
!$omp target update if(do_omp_offload) from(a1, x1, y1)
27132767
do l = 1, num_fields
27142768
write(text, '(i3.3)') l
27152769
call compare_checksums(a1(isd:ied, jsd:jed, :,l),a2(isd:ied, jsd:jed, :,l),type//' CENTER '//text)
@@ -2791,19 +2845,34 @@ subroutine test_group_update( type )
27912845
call mpp_create_group_update(group_update, a1(:,:,:,:), domain)
27922846

27932847
a1 = 0; x1 = 0; y1 = 0
2848+
!$omp target update if(do_omp_offload) to(a1)
2849+
!$omp target teams loop collapse(4) if(do_omp_offload) default(shared) &
2850+
!$omp shared(num_fields, nz, jsc, jec, isc, iec, a1, base) map(to: base) map(from: a1)
27942851
do l = 1, num_fields
2795-
a1(isc:iec, jsc:jec, :,l) = base(isc:iec, jsc:jec, :) + l*1e3
2852+
do k=1,nz
2853+
do j=jsc,jec
2854+
do i=isc,iec
2855+
a1(i,j,k,l) = base(i,j,k) + l*1e3
2856+
enddo
2857+
enddo
2858+
enddo
27962859
enddo
2797-
a2 = a1; x2 = x1; y2 = y1
2860+
!$omp target teams loop collapse(4) if(do_omp_offload) default(shared) &
2861+
!$omp shared(num_fields, nz, jsm, jem, ism, iem, a1, a2) map(from: a2) map(to: a1)
2862+
do l=1,num_fields ; do k=1,nz ; do j=jsm,jem ; do i=ism,iem
2863+
a2(i,j,k,l) = a1(i,j,k,l)
2864+
enddo ; enddo ; enddo ; enddo
2865+
x2 = x1; y2 = y1
27982866
call mpp_clock_begin(id1)
2799-
call mpp_do_group_update(group_update, domain, a1(isc,jsc,1,1))
2867+
call mpp_do_group_update(group_update, domain, a1(isc,jsc,1,1), omp_offload=do_omp_offload)
28002868
call mpp_clock_end (id1)
28012869

28022870
call mpp_clock_begin(id2)
28032871
call mpp_update_domains( a2(:,:,:,:), domain )
28042872
call mpp_clock_end(id2)
28052873

28062874
!--- compare checksum
2875+
!$omp target update if(do_omp_offload) from(a1)
28072876
do l = 1, num_fields
28082877
write(text, '(i3.3)') l
28092878
call compare_checksums(a1(isd:ied, jsd:jed, :,l),a2(isd:ied, jsd:jed, :,l),type//' 4D CENTER '//text)
@@ -2825,7 +2894,7 @@ subroutine test_group_update( type )
28252894
type//' nonblock 4D CENTER '//text)
28262895
enddo
28272896

2828-
2897+
!$omp target exit data if(do_omp_offload) map(release: a1, x1, y1)
28292898

28302899
!--- test for BGRID.
28312900
deallocate(a1, x1, y1)
@@ -2839,21 +2908,40 @@ subroutine test_group_update( type )
28392908
allocate( x2(ism:iem+shift,jsm:jem+shift, nz, num_fields) )
28402909
allocate( y2(ism:iem+shift,jsm:jem+shift, nz, num_fields) )
28412910

2911+
!$omp target enter data map(alloc: a1, x1, y1) if(do_omp_offload)
2912+
28422913
do l =1, num_fields
28432914
call mpp_create_group_update(group_update, a1(:,:,:,l), domain, position=CORNER)
28442915
call mpp_create_group_update(group_update, x1(:,:,:,l), y1(:,:,:,l), domain, gridtype=BGRID_NE)
28452916
end do
28462917

28472918
do n = 1, num_iter
28482919
a1 = 0; x1 = 0; y1 = 0
2920+
!$omp target update if(do_omp_offload) to(a1, x1, y1)
28492921
do l = 1, num_fields
2850-
a1(isc:iec+shift,jsc:jec+shift,:,l) = base(isc:iec+shift,jsc:jec+shift,:) + l*1e3
2851-
x1(isc:iec+shift,jsc:jec+shift,:,l) = base(isc:iec+shift,jsc:jec+shift,:) + l*1e3 + 1e6
2852-
y1(isc:iec+shift,jsc:jec+shift,:,l) = base(isc:iec+shift,jsc:jec+shift,:) + l*1e3 + 2*1e6
2922+
!$omp target teams loop collapse(3) if(do_omp_offload) default(shared) &
2923+
!$omp shared(nz, base, a1, x1, y1, l, jsc, jec, shift, isc, iec) &
2924+
!$omp map(to: base) map(tofrom: a1, x1, y1)
2925+
do k=1,nz
2926+
do j=jsc,jec+shift
2927+
do i=isc,iec+shift
2928+
a1(i,j,k,l) = base(i,j,k) + l*1e3
2929+
x1(i,j,k,l) = base(i,j,k) + l*1e3 + 1e6
2930+
y1(i,j,k,l) = base(i,j,k) + l*1e3 + 2*1e6
2931+
enddo
2932+
enddo
2933+
enddo
28532934
enddo
2854-
a2 = a1; x2 = x1; y2 = y1
2935+
!$omp target teams loop collapse(4) if(do_omp_offload) default(shared) &
2936+
!$omp shared(num_fields, nz, jsm, jem, shift, ism, iem, a1, x1, y1, a2, x2, y2) &
2937+
!$omp map(from: a2, x2, y2) map(to: a1, x1, y1)
2938+
do l=1,num_fields ; do k=1,nz ; do j=jsm,jem+shift ; do i=ism,iem+shift
2939+
a2(i,j,k,l) = a1(i,j,k,l)
2940+
x2(i,j,k,l) = x1(i,j,k,l)
2941+
y2(i,j,k,l) = y1(i,j,k,l)
2942+
enddo ; enddo ; enddo ; enddo
28552943
call mpp_clock_begin(id1)
2856-
call mpp_do_group_update(group_update, domain, a1(isc,jsc,1,1))
2944+
call mpp_do_group_update(group_update, domain, a1(isc,jsc,1,1), omp_offload=do_omp_offload)
28572945
call mpp_clock_end (id1)
28582946

28592947
call mpp_clock_begin(id2)
@@ -2867,6 +2955,7 @@ subroutine test_group_update( type )
28672955

28682956
!--- compare checksum
28692957
if( n == num_iter ) then
2958+
!$omp target update if(do_omp_offload) from(a1,x1,y1)
28702959
do l = 1, num_fields
28712960
write(text, '(i3.3)') l
28722961
call compare_checksums(a1(isd:ied+shift,jsd:jed+shift,:,l),a2(isd:ied+shift,jsd:jed+shift,:,l),type//&
@@ -2901,7 +2990,7 @@ subroutine test_group_update( type )
29012990
enddo
29022991
endif
29032992
enddo
2904-
2993+
!$omp target exit data if(do_omp_offload) map(release: a1,x1,y1)
29052994
call mpp_clear_group_update(group_update)
29062995

29072996
!-----------------------------------------------------------------------------

‎test_fms/mpp/test_mpp_domains2.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
. ../test-lib.sh
3030

3131
# TODO: Enable this test once generalized indices work is complete
32-
SKIP_TESTS="test_mpp_domains2.12 test_mpp_domains2.16"
32+
SKIP_TESTS="test_mpp_domains2.12"
3333

3434
# TODO edge update, fails on non-blocking with gnu
3535
#SKIP_TESTS="$SKIP_TESTS $(basename $0 .sh).6"

0 commit comments

Comments
 (0)