Add threading to the filter_topo and orog_gsl codes (ufs-community#948)

Include threading in both programs. Update the "regional.gsl.gwd" regression test - which runs both programs - to run twice with varying thread counts. Fixes ufs-community#939.
GeorgeGayno-NOAA · May 20, 2024 · c6efbbd · c6efbbd
1 parent 7034628
commit c6efbbd
Show file tree

Hide file tree

Showing 12 changed files with 132 additions and 45 deletions.
diff --git a/reg_tests/grid_gen/driver.hera.sh b/reg_tests/grid_gen/driver.hera.sh
@@ -103,19 +103,26 @@ TEST5=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:07:00 -A $PROJECT_
       -o $LOG_FILE5 -e $LOG_FILE5 ./esg.regional.pct.cat.sh)
 
 #-----------------------------------------------------------------------------
-# Regional GSL gravity wave drag test.
+# Regional GSL gravity wave drag test. This test is run with varying
+# thread counts.
 #-----------------------------------------------------------------------------
 
+export nthreads=12
 LOG_FILE6=${LOG_FILE}06
-TEST6=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:07:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd \
+TEST6=$(sbatch --parsable --ntasks-per-node=12 --nodes=1 -t 0:07:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd.12 \
       -o $LOG_FILE6 -e $LOG_FILE6 ./regional.gsl.gwd.sh)
 
+export nthreads=24
+LOG_FILE7=${LOG_FILE}07
+TEST7=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:07:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd.24 \
+      -o $LOG_FILE7 -e $LOG_FILE7 ./regional.gsl.gwd.sh)
+
 #-----------------------------------------------------------------------------
 # Create summary log.
 #-----------------------------------------------------------------------------
 
 sbatch --nodes=1 -t 0:01:00 -A $PROJECT_CODE -J grid_summary -o $LOG_FILE -e $LOG_FILE \
-       --open-mode=append -q $QUEUE -d afterok:$TEST1:$TEST2:$TEST3:$TEST4:$TEST5:$TEST6 << EOF
+       --open-mode=append -q $QUEUE -d afterok:$TEST1:$TEST2:$TEST3:$TEST4:$TEST5:$TEST6:$TEST7 << EOF
 #!/bin/bash
 grep -a '<<<' ${LOG_FILE}*  > $SUM_FILE
 EOF
diff --git a/reg_tests/grid_gen/driver.hercules.sh b/reg_tests/grid_gen/driver.hercules.sh
@@ -100,19 +100,27 @@ TEST5=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:10:00 -A $PROJECT_
       -o $LOG_FILE5 -e $LOG_FILE5 ./esg.regional.pct.cat.sh)
 
 #-----------------------------------------------------------------------------
-# Regional grid with GSL gravity wave drag fields.
+# Regional grid with GSL gravity wave drag fields. Run with varying
+# thread counts.
 #-----------------------------------------------------------------------------
 
+export nthreads=12
 LOG_FILE6=${LOG_FILE}06
-TEST6=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:10:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd \
+TEST6=$(sbatch --parsable --ntasks-per-node=12 --nodes=1 -t 0:10:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd.12 \
       -o $LOG_FILE6 -e $LOG_FILE6 ./regional.gsl.gwd.sh)
 
+
+export nthreads=24
+LOG_FILE7=${LOG_FILE}07
+TEST7=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:10:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd.24 \
+      -o $LOG_FILE7 -e $LOG_FILE7 ./regional.gsl.gwd.sh)
+
 #-----------------------------------------------------------------------------
 # Create summary log.
 #-----------------------------------------------------------------------------
 
 sbatch --nodes=1 -t 0:01:00 -A $PROJECT_CODE -J grid_summary -o $LOG_FILE -e $LOG_FILE \
-       -q $QUEUE -d afterok:$TEST1:$TEST2:$TEST3:$TEST4:$TEST5:$TEST6 << EOF
+       -q $QUEUE -d afterok:$TEST1:$TEST2:$TEST3:$TEST4:$TEST5:$TEST6:$TEST7 << EOF
 #!/bin/bash
 grep -a '<<<' ${LOG_FILE}*  > $SUM_FILE
 EOF
diff --git a/reg_tests/grid_gen/driver.jet.sh b/reg_tests/grid_gen/driver.jet.sh
@@ -101,13 +101,20 @@ TEST5=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:07:00 -A $PROJECT_
       --partition=xjet -o $LOG_FILE5 -e $LOG_FILE5 ./esg.regional.pct.cat.sh)
 
 #-----------------------------------------------------------------------------
-# Regional GSL gravity wave drag.
+# Regional GSL gravity wave drag. Run with varying
+# thread counts.
 #-----------------------------------------------------------------------------
 
+export nthreads=12
 LOG_FILE6=${LOG_FILE}06
-TEST6=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:07:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd \
+TEST6=$(sbatch --parsable --ntasks-per-node=12 --nodes=1 -t 0:07:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd.12 \
       --partition=xjet -o $LOG_FILE6 -e $LOG_FILE6 ./regional.gsl.gwd.sh)
 
+export nthreads=24
+LOG_FILE7=${LOG_FILE}07
+TEST7=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:07:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd.24 \
+      --partition=xjet -o $LOG_FILE7 -e $LOG_FILE7 ./regional.gsl.gwd.sh)
+
 #-----------------------------------------------------------------------------
 # Create summary log.
 #-----------------------------------------------------------------------------

diff --git a/reg_tests/grid_gen/driver.orion.sh b/reg_tests/grid_gen/driver.orion.sh
@@ -99,19 +99,26 @@ TEST5=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:10:00 -A $PROJECT_
       -o $LOG_FILE5 -e $LOG_FILE5 ./esg.regional.pct.cat.sh)
 
 #-----------------------------------------------------------------------------
-# Regional grid with GSL gravity wave drag fields.
+# Regional grid with GSL gravity wave drag fields. Run with varying
+# thread counts.
 #-----------------------------------------------------------------------------
 
+export nthreads=12
 LOG_FILE6=${LOG_FILE}06
-TEST6=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:10:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd \
+TEST6=$(sbatch --parsable --ntasks-per-node=12 --nodes=1 -t 0:10:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd.12 \
       -o $LOG_FILE6 -e $LOG_FILE6 ./regional.gsl.gwd.sh)
 
+export nthreads=24
+LOG_FILE7=${LOG_FILE}07
+TEST7=$(sbatch --parsable --ntasks-per-node=24 --nodes=1 -t 0:10:00 -A $PROJECT_CODE -q $QUEUE -J reg.gsl.gwd.24 \
+      -o $LOG_FILE7 -e $LOG_FILE7 ./regional.gsl.gwd.sh)
+
 #-----------------------------------------------------------------------------
 # Create summary log.
 #-----------------------------------------------------------------------------
 
 sbatch --nodes=1 -t 0:01:00 -A $PROJECT_CODE -J grid_summary -o $LOG_FILE -e $LOG_FILE \
-       -q $QUEUE -d afterok:$TEST1:$TEST2:$TEST3:$TEST4:$TEST5:$TEST6 << EOF
+       -q $QUEUE -d afterok:$TEST1:$TEST2:$TEST3:$TEST4:$TEST5:$TEST6:$TEST7 << EOF
 #!/bin/bash
 grep -a '<<<' ${LOG_FILE}*  > $SUM_FILE
 EOF
diff --git a/reg_tests/grid_gen/driver.wcoss2.sh b/reg_tests/grid_gen/driver.wcoss2.sh
@@ -104,19 +104,26 @@ TEST5=$(qsub -V -o $LOG_FILE5 -e $LOG_FILE5 -q $QUEUE -A $PROJECT_CODE -l wallti
         -N esg.regional.pct.cat -l select=1:ncpus=30:mem=40GB $PWD/esg.regional.pct.cat.sh)
 
 #-----------------------------------------------------------------------------
-# Regional GSL gravity wave drag test.
+# Regional GSL gravity wave drag test.  Run with varying
+# thread counts.
 #-----------------------------------------------------------------------------
 
+export nthreads=15
 LOG_FILE6=${LOG_FILE}06
 TEST6=$(qsub -V -o $LOG_FILE6 -e $LOG_FILE6 -q $QUEUE -A $PROJECT_CODE -l walltime=00:07:00 \
-        -N rsg.gsl.gwd -l select=1:ncpus=30:mem=40GB $PWD/regional.gsl.gwd.sh)
+        -N reg.gsl.gwd.15 -l select=1:ncpus=15:mem=40GB $PWD/regional.gsl.gwd.sh)
+
+export nthreads=30
+LOG_FILE7=${LOG_FILE}07
+TEST7=$(qsub -V -o $LOG_FILE7 -e $LOG_FILE7 -q $QUEUE -A $PROJECT_CODE -l walltime=00:07:00 \
+        -N reg.gsl.gwd -l select=1:ncpus=30:mem=40GB $PWD/regional.gsl.gwd.sh)
 
 #-----------------------------------------------------------------------------
 # Create summary log.
 #-----------------------------------------------------------------------------
 
 qsub -V -o ${LOG_FILE} -e ${LOG_FILE} -q $QUEUE -A $PROJECT_CODE -l walltime=00:02:00 \
-        -N grid_summary -l select=1:ncpus=1:mem=100MB -W depend=afterok:$TEST1:$TEST2:$TEST3:$TEST4:$TEST5:$TEST6 << EOF
+        -N grid_summary -l select=1:ncpus=1:mem=100MB -W depend=afterok:$TEST1:$TEST2:$TEST3:$TEST4:$TEST5:$TEST6:$TEST7 << EOF
 #!/bin/bash
 cd ${this_dir}
 grep -a '<<<' ${LOG_FILE}* | grep -v echo > $SUM_FILE

diff --git a/reg_tests/grid_gen/regional.gsl.gwd.sh b/reg_tests/grid_gen/regional.gsl.gwd.sh
@@ -8,8 +8,11 @@
 
 set -x
 
-export TEMP_DIR=${WORK_DIR}/regional.gsl.gwd.work
-export out_dir=${WORK_DIR}/regional.gsl.gwd
+nthreads=${nthreads:-6}
+export OMP_NUM_THREADS=$nthreads
+
+export TEMP_DIR=${WORK_DIR}/regional.gsl.gwd.${nthreads}.work
+export out_dir=${WORK_DIR}/regional.gsl.gwd.${nthreads}
 
 export gtype=regional_esg
 export make_gsl_orog=true    # Create GSL gravity wave drag fields
@@ -34,7 +37,7 @@ $home_dir/ush/fv3gfs_driver_grid.sh
 iret=$?
 if [ $iret -ne 0 ]; then
   set +x
-  echo "<<< REGIONAL GSL GWD TEST FAILED. <<<"
+  echo "<<< REGIONAL ${nthreads} THREAD GSL GWD TEST FAILED. <<<"
   exit $iret
 fi
 
@@ -61,12 +64,12 @@ done
 
 set +x
 if [ $test_failed -ne 0 ]; then
-  echo "<<< REGIONAL GSL GWD TEST FAILED. >>>"
+  echo "<<< REGIONAL ${nthreads} THREAD GSL GWD TEST FAILED. >>>"
   if [ "$UPDATE_BASELINE" = "TRUE" ]; then
     $home_dir/reg_tests/update_baseline.sh "${HOMEreg}/.." "regional.gsl.gwd" $commit_num
   fi
 else
-  echo "<<< REGIONAL GSL GWD TEST PASSED. >>>"
+  echo "<<< REGIONAL ${nthreads} THREAD GSL GWD TEST PASSED. >>>"
 fi
 
 exit 0
diff --git a/sorc/grid_tools.fd/filter_topo.fd/CMakeLists.txt b/sorc/grid_tools.fd/filter_topo.fd/CMakeLists.txt
@@ -24,6 +24,10 @@ target_link_libraries(
   PUBLIC
   NetCDF::NetCDF_Fortran)
 
+if(OpenMP_Fortran_FOUND)
+  target_link_libraries(filter_topo_lib PUBLIC OpenMP::OpenMP_Fortran)
+endif()
+
 target_link_libraries(${exe_name} PRIVATE filter_topo_lib)
 
 install(TARGETS ${exe_name})
diff --git a/sorc/grid_tools.fd/filter_topo.fd/filter_topo.F90 b/sorc/grid_tools.fd/filter_topo.fd/filter_topo.F90
@@ -9,6 +9,7 @@
 !! @author Zhi Liang (GFDL) who packaged it into a standalone application.
 program filter_topo
 
+  use omp_lib
   use utils
 
   implicit none
@@ -34,7 +35,7 @@ program filter_topo
   real:: peak_fac   ! overshoot factor for the mountain peak
   real:: max_slope  ! max allowable terrain slope: 1 --> 45 deg
 
-  integer :: n_del2_weak
+  integer :: n_del2_weak, tid, nthreads
 
   integer :: ntiles = 0
 
@@ -50,6 +51,17 @@ program filter_topo
   integer           :: is,ie,js,je,isd,ied,jsd,jed
   integer,parameter :: ng = 3
   integer           :: nx, ny, npx, npy, nx_nest, ny_nest, npx_nest, npy_nest, is_nest, ie_nest, js_nest, je_nest, isd_nest, ied_nest, jsd_nest, jed_nest
+
+!$OMP PARALLEL PRIVATE(TID)
+  tid = omp_get_thread_num()
+  if (tid == 0) then
+    nthreads = omp_get_num_threads()
+    print*
+    print*,'- BEGIN EXECUTION WITH NUMBER OF THREADS = ',nthreads
+    print*
+  endif
+!$OMP END PARALLEL
+
   !--- read namelist
   call read_namelist()
 
@@ -69,6 +81,9 @@ program filter_topo
   !--- write out the data
   call write_topo_file(is,ie,js,je,ntiles,oro(is:ie,js:je,:),regional )
 
+  print*
+  print*,'- NORMAL TERMINATION.'
+
 contains
 
   !> ???
@@ -760,6 +775,7 @@ subroutine read_grid_file(regional)
     geolat_t(:,:,:) = -1.e25
 
     do t = 1, ntiles
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,G1,G2,G3,G4,G5)
        do j=js,je ; do i=is,ie
           g1(1) = geolon_c(i,j,t);     g1(2) = geolat_c(i,j,t)
           g2(1) = geolon_c(i+1,j,t);   g2(2) = geolat_c(i+1,j,t)
@@ -769,8 +785,8 @@ subroutine read_grid_file(regional)
           geolon_t(i,j,t) = g5(1)
           geolat_t(i,j,t) = g5(2)
        enddo ; enddo
+!$OMP END PARALLEL DO
     enddo
-
 
     if( .not. regional ) then
       call fill_cubic_grid_halo(geolon_t, geolon_t, ng, 0, 0, 1, 1)
@@ -783,15 +799,18 @@ subroutine read_grid_file(regional)
     allocate(dx(isd:ied,jsd:jed+1,ntiles))
     allocate(dy(isd:ied+1,jsd:jed,ntiles))
     do t = 1, ntiles
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,G1,G2)
        do j = js, je+1 ; do i = is, ie
           g1(1) = geolon_c(i  ,j,t)
           g1(2) = geolat_c(i  ,j,t)
           g2(1) = geolon_c(i+1,j,t)
           g2(2) = geolat_c(i+1,j,t)
           dx(i,j,t) = great_circle_dist( g2, g1, radius )
        enddo ; enddo
+!$OMP END PARALLEL DO
     enddo
     do t = 1, ntiles
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,G1,G2)
        do j = js, je
           do i = is, ie+1
              g1(1) = geolon_c(i,j,  t)
@@ -801,6 +820,7 @@ subroutine read_grid_file(regional)
              dy(i,j,t) = great_circle_dist( g2, g1, radius )
           enddo
        enddo
+!$OMP END PARALLEL DO
     enddo
 
     if( .not. regional ) then
@@ -831,6 +851,7 @@ subroutine read_grid_file(regional)
     allocate(dxa(isd:ied,jsd:jed,ntiles))
     allocate(dya(isd:ied,jsd:jed,ntiles))
     do t = 1, ntiles
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,G1,G2,G3,G4)
        do j=js,je ; do i=is,ie
           g1(1) = geolon_c(i,j,t); g1(2) = geolat_c(i,j,t)
           g2(1) = geolon_c(i,j+1,t); g2(2) = geolat_c(i,j+1,t)
@@ -847,6 +868,7 @@ subroutine read_grid_file(regional)
           call mid_pt_sphere(g1, g2, g4)
           dya(i,j,t) = great_circle_dist( g4, g3, radius )
        enddo; enddo
+!$OMP END PARALLEL DO
     enddo
 
     if( .not.regional ) then
@@ -860,6 +882,8 @@ subroutine read_grid_file(regional)
     allocate(dxc(isd:ied+1,jsd:jed,ntiles))
     allocate(dyc(isd:ied,jsd:jed+1,ntiles))
     do t = 1, ntiles
+
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,G1,G2)
        do j=jsd,jed
           do i=isd+1,ied
              g1(1) = geolon_c(i,j,t); g1(2) = geolat_c(i,j,t)
@@ -869,23 +893,30 @@ subroutine read_grid_file(regional)
           dxc(isd,j,t)   = dxc(isd+1,j,t)
           dxc(ied+1,j,t) = dxc(ied,j,t)
        enddo
+!$OMP END PARALLEL DO
 
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,G1,G2)
        do j=jsd+1,jed
           do i=isd,ied
              g1(1) = geolon_c(i,j,t); g1(2) = geolat_c(i,j,t)
              g2(1) = geolon_c(i,j-1,t); g2(2) = geolat_c(i,j-1,t)
              dyc(i,j,t) = great_circle_dist(g1, g2, radius)
           enddo
        enddo
+!$OMP END PARALLEL DO
+
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I)
        do i=isd,ied
           dyc(i,jsd,t)   = dyc(i,jsd+1,t)
           dyc(i,jed+1,t) = dyc(i,jed,t)
        end do
+!$OMP END PARALLEL DO
     enddo
 
     !--- compute area
     allocate(area(isd:ied,jsd:jed,ntiles))
     do t = 1, ntiles
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,p_lL,p_uL,p_lr,p_uR)
        do j=js,je
           do i=is,ie
              p_lL(1) = geolon_c(i  ,j  ,t) ; p_lL(2) = geolat_c(i  ,j  ,t)
@@ -897,6 +928,7 @@ subroutine read_grid_file(regional)
              area(i,j,t) = get_area(p_lL, p_uL, p_lR, p_uR, radius)
           enddo
        enddo
+!$OMP END PARALLEL DO
     enddo
 
     if( .not.regional ) then
@@ -918,13 +950,18 @@ subroutine read_grid_file(regional)
     !     |       |
     !     6---2---7
     do t = 1, ntiles
+
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,G1)
        do j=js,je+1
           do i = is,ie+1
              g1(1) = geolon_c(i,j,t)
              g1(2) = geolat_c(i,j,t)
              call latlon2xyz(g1, grid3(:,i,j))
           enddo
        enddo
+!$OMP END PARALLEL DO
+
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J,G1,P1,P3)
        do j=js,je
           do i=is,ie
              g1(1) = geolon_t(i,j,t); g1(2) = geolat_t(i,j,t)
@@ -939,13 +976,16 @@ subroutine read_grid_file(regional)
              cos_sg(4,i,j) = cos_angle( p1, grid3(1,i,j+1), p3 )
           enddo
        enddo
+!$OMP END PARALLEL DO
 
        do ip=1,4
+!$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(I,J)
           do j=js,je
              do i=is,ie
                 sin_sg(ip,i,j,t) = min(1.0, sqrt( max(0., 1.-cos_sg(ip,i,j)**2) ) )
              enddo
           enddo
+!$OMP END PARALLEL DO
        enddo
     enddo
 

diff --git a/sorc/orog_mask_tools.fd/orog_gsl.fd/CMakeLists.txt b/sorc/orog_mask_tools.fd/orog_gsl.fd/CMakeLists.txt
@@ -26,6 +26,10 @@ target_link_libraries(
   PUBLIC
   NetCDF::NetCDF_Fortran)
 
+if(OpenMP_Fortran_FOUND)
+  target_link_libraries(orog_gsl_lib PUBLIC OpenMP::OpenMP_Fortran)
+endif()
+
 target_link_libraries(${exe_name} PRIVATE orog_gsl_lib)
 
 install(TARGETS ${exe_name})