[Draft] Consolidating OpenACC device-host memory transfers

src/core_atmosphere/dynamics/mpas_atm_boundaries.F

@@ -287,7 +287,7 @@ end subroutine mpas_atm_update_bdy_tend
     !>   tend_scalars(1,:,:) = mpas_atm_get_bdy_tend(clock, domain % blocklist, nVertLevels, nCells, 'qv', 0.0_RKIND)
     !
     !-----------------------------------------------------------------------
-    function mpas_atm_get_bdy_tend(clock, block, vertDim, horizDim, field, delta_t) result(return_tend)
+    subroutine mpas_atm_get_bdy_tend(clock, block, vertDim, horizDim, field, delta_t, return_tend)
 
         implicit none
 
@@ -296,31 +296,49 @@ function mpas_atm_get_bdy_tend(clock, block, vertDim, horizDim, field, delta_t)
         integer, intent(in) :: vertDim, horizDim
         character(len=*), intent(in) :: field
         real (kind=RKIND), intent(in) :: delta_t
-
-        real (kind=RKIND), dimension(vertDim,horizDim+1) :: return_tend
+        real (kind=RKIND), dimension(vertDim,horizDim+1), intent(out) :: return_tend
 
         type (mpas_pool_type), pointer :: lbc
-        integer, pointer :: idx
+        integer, pointer :: idx_ptr
         real (kind=RKIND), dimension(:,:), pointer :: tend
         real (kind=RKIND), dimension(:,:,:), pointer :: tend_scalars
-        integer :: ierr
+        integer :: idx, i, j
 
 
         call mpas_pool_get_subpool(block % structs, 'lbc', lbc)
 
         nullify(tend)
         call mpas_pool_get_array(lbc, 'lbc_'//trim(field), tend, 1)
-
+        
         if (associated(tend)) then
-            return_tend(:,:) = tend(:,:)
         else
             call mpas_pool_get_array(lbc, 'lbc_scalars', tend_scalars, 1)
-            call mpas_pool_get_dimension(lbc, 'index_'//trim(field), idx)
 
-            return_tend(:,:) = tend_scalars(idx,:,:)
+            ! Ensure the integer pointed to by idx_ptr is copied to the gpu device
+            call mpas_pool_get_dimension(lbc, 'index_'//trim(field), idx_ptr)
+            idx = idx_ptr
+        end if
+
+        !$acc parallel default(present)
+        if (associated(tend)) then
+            !$acc loop gang vector collapse(2)
+            do j=1,horizDim+1
+                do i=1,vertDim
+                    return_tend(i,j) = tend(i,j)
+                end do
+            end do
+        else
+            !$acc loop gang vector collapse(2)
+            do j=1,horizDim+1
+                do i=1,vertDim
+                    return_tend(i,j) = tend_scalars(idx,i,j)
+                end do
+            end do
         end if
+        !$acc end parallel
+
 
-    end function mpas_atm_get_bdy_tend
+    end subroutine mpas_atm_get_bdy_tend
 
 
     !***********************************************************************
@@ -356,10 +374,11 @@ end function mpas_atm_get_bdy_tend
     !>   scalars(1,:,:) = mpas_atm_get_bdy_state(clock, domain % blocklist, nVertLevels, nCells, 'qv', 0.0_RKIND)
     !
     !-----------------------------------------------------------------------
-    function mpas_atm_get_bdy_state_2d(clock, block, vertDim, horizDim, field, delta_t) result(return_state)
+    subroutine mpas_atm_get_bdy_state_2d(clock, block, vertDim, horizDim, field, delta_t, return_state)
 
         use mpas_pool_routines, only : mpas_pool_get_error_level, mpas_pool_set_error_level
         use mpas_derived_types, only : MPAS_POOL_SILENT
+        use mpas_log, only : mpas_log_write
 
         implicit none
 
@@ -368,8 +387,7 @@ function mpas_atm_get_bdy_state_2d(clock, block, vertDim, horizDim, field, delta
         integer, intent(in) :: vertDim, horizDim
         character(len=*), intent(in) :: field
         real (kind=RKIND), intent(in) :: delta_t
-
-        real (kind=RKIND), dimension(vertDim,horizDim+1) :: return_state
+        real (kind=RKIND), dimension(vertDim,horizDim+1), intent(out) :: return_state
 
         type (mpas_pool_type), pointer :: lbc
         integer, pointer :: idx_ptr
@@ -420,10 +438,6 @@ function mpas_atm_get_bdy_state_2d(clock, block, vertDim, horizDim, field, delta
         ! query the field as a scalar constituent
         !
         if (associated(tend) .and. associated(state)) then
-            MPAS_ACC_TIMER_START('mpas_atm_get_bdy_state_2d [ACC_data_xfer]')
-            !$acc enter data create(return_state) &
-            !$acc            copyin(tend, state)
-            MPAS_ACC_TIMER_STOP('mpas_atm_get_bdy_state_2d [ACC_data_xfer]')
 
             !$acc parallel default(present)
             !$acc loop gang vector collapse(2)
@@ -434,22 +448,13 @@ function mpas_atm_get_bdy_state_2d(clock, block, vertDim, horizDim, field, delta
             end do
             !$acc end parallel
 
-            MPAS_ACC_TIMER_START('mpas_atm_get_bdy_state_2d [ACC_data_xfer]')
-            !$acc exit data copyout(return_state) &
-            !$acc           delete(tend, state)
-            MPAS_ACC_TIMER_STOP('mpas_atm_get_bdy_state_2d [ACC_data_xfer]')
         else
             call mpas_pool_get_array(lbc, 'lbc_scalars', tend_scalars, 1)
             call mpas_pool_get_array(lbc, 'lbc_scalars', state_scalars, 2)
             call mpas_pool_get_dimension(lbc, 'index_'//trim(field), idx_ptr)
 
             idx=idx_ptr ! Avoid non-array pointer for OpenACC
 
-            MPAS_ACC_TIMER_START('mpas_atm_get_bdy_state_2d [ACC_data_xfer]')
-            !$acc enter data create(return_state) &
-            !$acc            copyin(tend_scalars, state_scalars)
-            MPAS_ACC_TIMER_STOP('mpas_atm_get_bdy_state_2d [ACC_data_xfer]')
-
             !$acc parallel default(present)
             !$acc loop gang vector collapse(2)
             do i=1, horizDim+1
@@ -459,13 +464,9 @@ function mpas_atm_get_bdy_state_2d(clock, block, vertDim, horizDim, field, delta
             end do
             !$acc end parallel
 
-            MPAS_ACC_TIMER_START('mpas_atm_get_bdy_state_2d [ACC_data_xfer]')
-            !$acc exit data copyout(return_state) &
-            !$acc           delete(tend_scalars, state_scalars)
-            MPAS_ACC_TIMER_STOP('mpas_atm_get_bdy_state_2d [ACC_data_xfer]')
         end if
 
-    end function mpas_atm_get_bdy_state_2d
+    end subroutine mpas_atm_get_bdy_state_2d
 
 
     !***********************************************************************
@@ -498,7 +499,7 @@ end function mpas_atm_get_bdy_state_2d
     !>                       num_scalars, nVertLevels, nCells, 'scalars', 0.0_RKIND)
     !
     !-----------------------------------------------------------------------
-    function mpas_atm_get_bdy_state_3d(clock, block, innerDim, vertDim, horizDim, field, delta_t) result(return_state)
+    subroutine mpas_atm_get_bdy_state_3d(clock, block, innerDim, vertDim, horizDim, field, delta_t, return_state)
 
         use mpas_pool_routines, only : mpas_pool_get_error_level, mpas_pool_set_error_level
         use mpas_derived_types, only : MPAS_POOL_SILENT
@@ -510,8 +511,7 @@ function mpas_atm_get_bdy_state_3d(clock, block, innerDim, vertDim, horizDim, fi
         integer, intent(in) :: innerDim, vertDim, horizDim
         character(len=*), intent(in) :: field
         real (kind=RKIND), intent(in) :: delta_t
-
-        real (kind=RKIND), dimension(innerDim,vertDim,horizDim+1) :: return_state
+        real (kind=RKIND), dimension(innerDim,vertDim,horizDim+1), intent(out) :: return_state
 
         type (mpas_pool_type), pointer :: lbc
         real (kind=RKIND), dimension(:,:,:), pointer :: tend
@@ -543,11 +543,6 @@ function mpas_atm_get_bdy_state_3d(clock, block, innerDim, vertDim, horizDim, fi
         call mpas_pool_get_array(lbc, 'lbc_'//trim(field), tend, 1)
         call mpas_pool_get_array(lbc, 'lbc_'//trim(field), state, 2)
 
-        MPAS_ACC_TIMER_START('mpas_atm_get_bdy_state_3d [ACC_data_xfer]')
-        !$acc enter data create(return_state) &
-        !$acc            copyin(tend, state)
-        MPAS_ACC_TIMER_STOP('mpas_atm_get_bdy_state_3d [ACC_data_xfer]')
-
         !$acc parallel default(present)
         !$acc loop gang vector collapse(3)
         do i=1, horizDim+1
@@ -559,12 +554,7 @@ function mpas_atm_get_bdy_state_3d(clock, block, innerDim, vertDim, horizDim, fi
         end do
         !$acc end parallel
 
-        MPAS_ACC_TIMER_START('mpas_atm_get_bdy_state_3d [ACC_data_xfer]')
-        !$acc exit data copyout(return_state) &
-        !$acc           delete(tend, state)
-        MPAS_ACC_TIMER_STOP('mpas_atm_get_bdy_state_3d [ACC_data_xfer]')
-
-    end function mpas_atm_get_bdy_state_3d
+    end subroutine mpas_atm_get_bdy_state_3d
 
 
     !***********************************************************************

src/core_atmosphere/dynamics/mpas_atm_iau.F

@@ -137,6 +137,7 @@ 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)
@@ -149,6 +150,7 @@ 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)
 
       call mpas_pool_get_array(tend_iau, 'theta',     theta_amb)
       call mpas_pool_get_array(tend_iau, 'rho',         rho_amb)