Moving some OpenACC data movements to subroutines

Author: abishekg7

src/core_atmosphere/dynamics/mpas_atm_iau.F

@@ -5,6 +5,15 @@
 ! Additional copyright and license information can be found in the LICENSE file
 ! distributed with this code, or at http://mpas-dev.github.com/license.html
 !
+
+ #ifdef MPAS_OPENACC
+ #define MPAS_ACC_TIMER_START(X) call mpas_timer_start(X)
+ #define MPAS_ACC_TIMER_STOP(X) call mpas_timer_stop(X)
+ #else
+ #define MPAS_ACC_TIMER_START(X)
+ #define MPAS_ACC_TIMER_STOP(X)
+ #endif
+
 module mpas_atm_iau
 
  use mpas_derived_types
@@ -15,17 +24,7 @@ module mpas_atm_iau
  use mpas_log, only : mpas_log_write
  use mpas_timer
 
- !public :: atm_compute_iau_coef, atm_add_tend_anal_incr  
-      
-
- #ifdef MPAS_OPENACC
- #define MPAS_ACC_TIMER_START(X) call mpas_timer_start(X)
- #define MPAS_ACC_TIMER_STOP(X) call mpas_timer_stop(X)
- #else
- #define MPAS_ACC_TIMER_START(X)
- #define MPAS_ACC_TIMER_STOP(X)
- #endif
- 
+ !public :: atm_compute_iau_coef, atm_add_tend_anal_incr    
 
  contains
 
@@ -87,6 +86,39 @@ real (kind=RKIND) function atm_iau_coef(configs, itimestep, dt) result(wgt_iau)
       end if  
 
  end function atm_iau_coef
+
+!==================================================================================================
+ subroutine update_d2h_pre_add_tend_anal_incr(configs,structs)
+!==================================================================================================
+  
+      implicit none
+
+      type (mpas_pool_type), intent(in) :: configs
+      type (mpas_pool_type), intent(inout) :: structs
+
+      type (mpas_pool_type), pointer :: tend
+      type (mpas_pool_type), pointer :: state
+      type (mpas_pool_type), pointer :: diag
+          
+      real (kind=RKIND), dimension(:,:), pointer :: rho_edge, rho_zz, theta_m
+      real(kind=RKIND),dimension(:,:,:), pointer :: scalars, tend_scalars
+
+      call mpas_pool_get_subpool(structs, 'tend', tend)
+      call mpas_pool_get_subpool(structs, 'state', state)
+      call mpas_pool_get_subpool(structs, 'diag', diag)
+
+      MPAS_ACC_TIMER_START('atm_srk3: physics ACC_data_xfer')
+      call mpas_pool_get_array(state, 'theta_m', theta_m, 1)
+      call mpas_pool_get_array(state, 'scalars', scalars, 1)
+      call mpas_pool_get_array(state, 'rho_zz',   rho_zz, 2) 
+      call mpas_pool_get_array(diag , 'rho_edge', rho_edge)
+      !$acc update self(theta_m, scalars, rho_zz, rho_edge)
+
+      call mpas_pool_get_array(tend, 'scalars_tend', tend_scalars)
+      !$acc update self(tend_scalars)
+      MPAS_ACC_TIMER_STOP('atm_srk3: physics ACC_data_xfer')
+
+ end subroutine update_d2h_pre_add_tend_anal_incr
 
 !==================================================================================================
  subroutine atm_add_tend_anal_incr (configs, structs, itimestep, dt, tend_ru, tend_rtheta, tend_rho)
@@ -148,7 +180,6 @@ subroutine atm_add_tend_anal_incr (configs, structs, itimestep, dt, tend_ru, ten
       call mpas_pool_get_array(state, 'scalars', scalars, 1)
       call mpas_pool_get_array(state, 'rho_zz',   rho_zz, 2) 
       call mpas_pool_get_array(diag , 'rho_edge', rho_edge)
-      !$acc update self(theta_m, scalars, rho_zz, rho_edge)
 
       call mpas_pool_get_dimension(state, 'moist_start', moist_start)
       call mpas_pool_get_dimension(state, 'moist_end', moist_end)
@@ -161,8 +192,6 @@ subroutine atm_add_tend_anal_incr (configs, structs, itimestep, dt, tend_ru, ten
 !      call mpas_pool_get_array(tend, 'rho_zz',  tend_rho)
 !      call mpas_pool_get_array(tend, 'theta_m', tend_theta)
       call mpas_pool_get_array(tend, 'scalars_tend', tend_scalars)
-      !$acc update self(tend_scalars)
-      MPAS_ACC_TIMER_STOP('atm_srk3: physics ACC_data_xfer')
 
       call mpas_pool_get_array(tend_iau, 'theta',     theta_amb)
       call mpas_pool_get_array(tend_iau, 'rho',         rho_amb)

src/core_atmosphere/dynamics/mpas_atm_time_integration.F

@@ -29,6 +29,7 @@ module atm_time_integration
 
 #ifdef DO_PHYSICS
    use mpas_atmphys_driver_microphysics
+   use mpas_atmphys_interface, only: update_d2h_pre_microphysics, update_h2d_post_microphysics
    use mpas_atmphys_todynamics
    use mpas_atmphys_utilities
 #endif
@@ -1985,6 +1986,7 @@ subroutine atm_srk3(domain, dt, itimestep, exchange_halo_group)
       call mpas_timer_stop('atm_compute_moist_coefficients')
 
 #ifdef DO_PHYSICS
+      call update_d2h_pre_physics_get_tend(block % configs, state, diag, tend)
       call mpas_timer_start('physics_get_tend')
       rk_step = 1
       dynamics_substep = 1
@@ -1993,6 +1995,7 @@ subroutine atm_srk3(domain, dt, itimestep, exchange_halo_group)
                              tend_ru_physics, tend_rtheta_physics, tend_rho_physics, &
                              exchange_halo_group )
       call mpas_timer_stop('physics_get_tend')
+      call update_h2d_post_physics_get_tend(block % configs, state, diag, tend)
 #else
 #ifndef MPAS_CAM_DYCORE
       !
@@ -2008,6 +2011,7 @@ subroutine atm_srk3(domain, dt, itimestep, exchange_halo_group)
       ! IAU - Incremental Analysis Update
       !
       if (trim(config_IAU_option) /= 'off') then
+         call update_d2h_pre_add_tend_anal_incr(block % configs, block % structs)
          call atm_add_tend_anal_incr(block % configs, block % structs, itimestep, dt, &
                                      tend_ru_physics, tend_rtheta_physics, tend_rho_physics)
       end if
@@ -2614,7 +2618,6 @@ subroutine atm_srk3(domain, dt, itimestep, exchange_halo_group)
          !$acc update self(rthdynten)
          MPAS_ACC_TIMER_STOP('atm_srk3: physics ACC_data_xfer')
 
-
          !NOTE: The calculation of the tendency due to horizontal and vertical advection for the water vapor mixing ratio
          !requires that the subroutine atm_advance_scalars_mono was called on the third Runge Kutta step, so that a halo
          !update for the scalars at time_levs(1) is applied. A halo update for the scalars at time_levs(2) is done above. 
@@ -2643,6 +2646,7 @@ subroutine atm_srk3(domain, dt, itimestep, exchange_halo_group)
       MPAS_ACC_TIMER_STOP('atm_srk3: physics ACC_data_xfer')
       !call microphysics schemes:
       if (trim(config_microp_scheme) /= 'off')  then
+        call update_d2h_pre_microphysics( block % configs, state, diag, 2)
         call mpas_timer_start('microphysics')
 !$OMP PARALLEL DO
         do thread=1,nThreads
@@ -2651,6 +2655,7 @@ subroutine atm_srk3(domain, dt, itimestep, exchange_halo_group)
         end do
 !$OMP END PARALLEL DO
         call mpas_timer_stop('microphysics')
+        call update_h2d_post_microphysics( block % configs, state, diag, tend, 2)
       end if
 
       !

src/core_atmosphere/physics/mpas_atmphys_interface.F

@@ -6,13 +6,6 @@
 ! distributed with this code, or at http://mpas-dev.github.com/license.html
 !
 !=================================================================================================================
- module mpas_atmphys_interface
- use mpas_kind_types
- use mpas_pool_routines
-
- use mpas_atmphys_constants
- use mpas_atmphys_vars
- use mpas_timer
 
 #ifdef MPAS_OPENACC
 #define MPAS_ACC_TIMER_START(X) call mpas_timer_start(X)
@@ -22,6 +15,13 @@ module mpas_atmphys_interface
 #define MPAS_ACC_TIMER_STOP(X)
 #endif
 
+ module mpas_atmphys_interface
+ use mpas_kind_types
+ use mpas_pool_routines
+
+ use mpas_atmphys_constants
+ use mpas_atmphys_vars
+ use mpas_timer
 
  implicit none
  private
@@ -555,6 +555,40 @@ subroutine MPAS_to_physics(configs,mesh,state,time_lev,diag,diag_physics,its,ite
 
  end subroutine MPAS_to_physics
 
+!=================================================================================================================
+ subroutine update_d2h_pre_microphysics(configs,state,diag,time_lev)
+!=================================================================================================================
+
+!input variables:
+ type(mpas_pool_type),intent(in):: configs
+ type(mpas_pool_type),intent(in):: state
+ type(mpas_pool_type),intent(in):: diag
+
+ integer:: time_lev
+
+!local pointers:
+ real(kind=RKIND),dimension(:,:),pointer  :: exner,pressure_b,w
+ real(kind=RKIND),dimension(:,:),pointer  :: rho_zz,theta_m,pressure_p
+ real(kind=RKIND),dimension(:,:,:),pointer:: scalars
+
+
+ MPAS_ACC_TIMER_START('update_d2h_pre_microphysics [ACC_data_xfer]')
+ call mpas_pool_get_array(diag,'exner'        ,exner     )
+ call mpas_pool_get_array(diag,'pressure_base',pressure_b)
+ call mpas_pool_get_array(diag,'pressure_p'   ,pressure_p)
+
+ call mpas_pool_get_array(state,'rho_zz' ,rho_zz ,time_lev)
+ call mpas_pool_get_array(state,'theta_m',theta_m,time_lev)
+ call mpas_pool_get_array(state,'w'      ,w      ,time_lev)
+ !$acc update host(exner, pressure_b, pressure_p, rho_zz, theta_m, w)
+
+ call mpas_pool_get_array(state,'scalars',scalars,time_lev)
+ !$acc update host(scalars)
+
+ MPAS_ACC_TIMER_STOP('update_d2h_pre_microphysics [ACC_data_xfer]')
+
+end subroutine update_d2h_pre_microphysics
+
 !=================================================================================================================
  subroutine microphysics_from_MPAS(configs,mesh,state,time_lev,diag,diag_physics,tend_physics,its,ite)
 !=================================================================================================================
@@ -598,22 +632,18 @@ subroutine microphysics_from_MPAS(configs,mesh,state,time_lev,diag,diag_physics,
  call mpas_pool_get_array(mesh,'zgrid',zgrid)
  call mpas_pool_get_array(mesh,'zz'   ,zz   )
 
- MPAS_ACC_TIMER_START('atm_srk3: physics ACC_data_xfer')
  call mpas_pool_get_array(diag,'exner'        ,exner     )
  call mpas_pool_get_array(diag,'pressure_base',pressure_b)
  call mpas_pool_get_array(diag,'pressure_p'   ,pressure_p)
 
  call mpas_pool_get_array(state,'rho_zz' ,rho_zz ,time_lev)
  call mpas_pool_get_array(state,'theta_m',theta_m,time_lev)
  call mpas_pool_get_array(state,'w'      ,w      ,time_lev)
- !$acc update host(exner, pressure_b, pressure_p, rho_zz, theta_m, w)
 
  call mpas_pool_get_dimension(state,'index_qv',index_qv)
  call mpas_pool_get_dimension(state,'index_qc',index_qc)
  call mpas_pool_get_dimension(state,'index_qr',index_qr)
- call mpas_pool_get_array(state,'scalars',scalars,time_lev)
- !$acc update host(scalars)
- MPAS_ACC_TIMER_STOP('atm_srk3: physics ACC_data_xfer')
+ call mpas_pool_get_array(state,'scalars',scalars,time_lev) 
  qv => scalars(index_qv,:,:)
  qc => scalars(index_qc,:,:)
  qr => scalars(index_qr,:,:)
@@ -1054,13 +1084,49 @@ subroutine microphysics_to_MPAS(configs,mesh,state,time_lev,diag,diag_physics,te
     case default
  end select mp_tend_select
 
- MPAS_ACC_TIMER_START('atm_srk3: physics ACC_data_xfer')
+ end subroutine microphysics_to_MPAS
+
+ !=================================================================================================================
+ subroutine update_h2d_post_microphysics(configs,state,diag,tend,time_lev)
+!=================================================================================================================
+
+!input variables:
+ type(mpas_pool_type),intent(in):: configs
+ type(mpas_pool_type),intent(in):: state
+ type(mpas_pool_type),intent(in):: diag
+ type(mpas_pool_type),intent(inout):: tend
+
+
+ integer:: time_lev
+
+!local pointers:
+ real(kind=RKIND),dimension(:,:),pointer  :: exner,exner_b,pressure_b,rtheta_p,rtheta_b
+ real(kind=RKIND),dimension(:,:),pointer  :: rho_zz,theta_m,pressure_p
+ real(kind=RKIND),dimension(:,:,:),pointer:: scalars
+ real(kind=RKIND),dimension(:,:),pointer  :: rt_diabatic_tend
+
+ call mpas_pool_get_array(diag,'exner'        ,exner     )
+ call mpas_pool_get_array(diag,'exner_base'   ,exner_b   )
+ call mpas_pool_get_array(diag,'pressure_base',pressure_b)
+ call mpas_pool_get_array(diag,'pressure_p'   ,pressure_p)
+ call mpas_pool_get_array(diag,'rtheta_base'  ,rtheta_b  )
+ call mpas_pool_get_array(diag,'rtheta_p'     ,rtheta_p  )
+
+ call mpas_pool_get_array(state,'rho_zz' ,rho_zz ,time_lev)
+ call mpas_pool_get_array(state,'theta_m',theta_m,time_lev)
+
+ call mpas_pool_get_array(state,'scalars',scalars,time_lev)
+
+ call mpas_pool_get_array(tend,'rt_diabatic_tend',rt_diabatic_tend)
+
+
+ MPAS_ACC_TIMER_START('update_h2d_post_microphysics [ACC_data_xfer]')
  !$acc update device(exner, exner_b, pressure_b, pressure_p, rtheta_b)
  !$acc update device(rtheta_p, rho_zz, theta_m, scalars)
  !$acc update device(rt_diabatic_tend)
- MPAS_ACC_TIMER_STOP('atm_srk3: physics ACC_data_xfer')
+ MPAS_ACC_TIMER_STOP('update_h2d_post_microphysics [ACC_data_xfer]')
 
- end subroutine microphysics_to_MPAS
+end subroutine update_h2d_post_microphysics
 
 !=================================================================================================================
  end module mpas_atmphys_interface