Skip to content

Hires tracer optimizations #132

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.

Sign up for GitHub

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

Open
wmputman wants to merge 2 commits into
base: geos/main
Choose a base branch
from
Open

Hires tracer optimizations #132

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cf2453b
015e just updated the logfile output for Rayleigh damping timescales
wmputman Aug 31, 2026
8f6a2d6
important hires tracer optimizations
wmputman Sep 30, 2026
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
60 changes: 46 additions & 14 deletions model/dyn_core.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ module dyn_core_mod
use mpp_parameter_mod, only: CORNER
use fv_mp_mod, only: is_master
use fv_mp_mod, only: start_group_halo_update, complete_group_halo_update
use fv_mp_mod, only: group_halo_update_type
use fv_mp_mod, only: group_halo_update_type, mp_ireduce_max
use sw_core_mod, only: c_sw, d_sw, d2a2c_vect
use a2b_edge_mod, only: a2b_ord2, a2b_ord4
use nh_core_mod, only: Riem_Solver3, Riem_Solver_C, update_dz_c, update_dz_d, nest_halo_nh
Expand Down Expand Up @@ -162,9 +162,9 @@ module dyn_core_mod

subroutine dyn_core(npx, npy, npz, ng, sphum, nq, bdt, k_split, n_split, zvir, cp, akap, cappa, grav, hydrostatic, &
u, v, w, delz, pt, q, delp, pe, pk, phis, varflt, ws, omga, ptop, pfull, ua, va, &
dudt_rf, dvdt_rf, dwdt_rf, uc, vc, mfx, mfy, cx, cy, pkz, peln, q_con, ak, bk, dpx, &
dudt_rf, dvdt_rf, dwdt_rf, uc, vc, mfx, mfy, cx, cy, cmax_z, imax_req, q_split, pkz, peln, q_con, ak, bk, dpx, &
ks, gridstruct, flagstruct, neststruct, idiag, bd, domain, &
init_step, i_pack, end_step, diss_est,time_total)
init_step, i_pack, end_step, diss_est, time_total)

integer, intent(IN) :: npx
integer, intent(IN) :: npy
Expand Down Expand Up @@ -225,6 +225,10 @@ subroutine dyn_core(npx, npy, npz, ng, sphum, nq, bdt, k_split, n_split, zvir, c
real, intent(inout), dimension(bd%isd:bd%ied,bd%jsd:bd%jed,npz):: ua, va
real, intent(inout):: q_con(bd%isd:, bd%jsd:, 1:)

real, intent(out):: cmax_z(npz)
integer, intent(out):: imax_req
integer, intent(in):: q_split

! The Flux capacitors: accumulated Mass flux arrays
real(kind=8), intent(inout):: mfx(bd%is:bd%ie+1, bd%js:bd%je, npz)
real(kind=8), intent(inout):: mfy(bd%is:bd%ie , bd%js:bd%je+1, npz)
Expand All @@ -242,6 +246,7 @@ subroutine dyn_core(npx, npy, npz, ng, sphum, nq, bdt, k_split, n_split, zvir, c
real, allocatable, dimension(:,:,:):: pem, heat_source
! Auto 1D & 2D arrays:
real, dimension(bd%isd:bd%ied,bd%jsd:bd%jed):: ws3, z_rat
real:: cmax
real:: dp_ref(npz)
real:: zs(bd%isd:bd%ied,bd%jsd:bd%jed) !< surface height (m)
real:: p1d(bd%is:bd%ie)
Expand Down Expand Up @@ -393,8 +398,6 @@ subroutine dyn_core(npx, npy, npz, ng, sphum, nq, bdt, k_split, n_split, zvir, c
endif
endif



!-----------------------------------------------------
do it=1,n_split
!-----------------------------------------------------
Expand Down Expand Up @@ -826,12 +829,48 @@ subroutine dyn_core(npx, npy, npz, ng, sphum, nq, bdt, k_split, n_split, zvir, c
enddo
enddo
endif

enddo ! end openMP k-loop
if ( (.not. hydrostatic) .and. flagstruct%fv_debug ) &
call prt_mxm('W_dsw ', w, is, ie , js, je , ng, npz, 1., gridstruct%area_64, domain)

call timing_off('d_sw')


! Initialize the request handle so the later WAIT call doesn't crash
! if the reduction is never started.
if ( (it==n_split) .and. nq > 0 .and. .not. flagstruct%inline_q ) then

! --- 1. NON-BLOCKING CMAX REDUCTION START ---
! A. Compute and store local cmax_z for this processor
!$OMP parallel do default(none) shared(npz,is,ie,js,je,cx,cy,cmax_z,gridstruct) &
!$OMP private(i,j,k,cmax)
do k = 1, npz
! Compute local cmax_z (cx and cy are already hot in cache)
cmax_z(k) = 0.
do j=js,je
do i=is,ie
cmax = max(abs(cx(i,j,k)),abs(cy(i,j,k)))+1.-gridstruct%sin_sg(i,j,5)
cmax_z(k) = max( cmax_z(k), cmax )
enddo
enddo
enddo
if ( q_split == 0 ) then
! B. Initiate the non-blocking global reduction
call timing_on('COMM_TOTAL')
call timing_on('COMM_TRACER_MAX')
call mp_ireduce_max(cmax_z, npz, imax_req)
call timing_off('COMM_TRACER_MAX')
call timing_off('COMM_TOTAL')
endif

call timing_on('COMM_TOTAL')
call timing_on('COMM_TRACER')
call start_group_halo_update(i_pack(10), q, domain)
call timing_off('COMM_TRACER')
call timing_off('COMM_TOTAL')
endif

if( flagstruct%fill_dp ) call mix_dp(hydrostatic, w, delp, pt, npz, ak, bk, .false., flagstruct%fv_debug, bd)

if ( flagstruct%fv_debug ) &
Expand Down Expand Up @@ -1042,6 +1081,7 @@ subroutine dyn_core(npx, npy, npz, ng, sphum, nq, bdt, k_split, n_split, zvir, c

! *** Inline Rayleigh friction here
if( flagstruct%RF_fast .and. flagstruct%tau > 0. ) then
call timing_on('RAY_FAST')
dudt_rf = u(is:ie,js:je,:)
dvdt_rf = v(is:ie,js:je,:)
if (.not. hydrostatic) dwdt_rf = w(is:ie,js:je,:)
Expand All @@ -1050,6 +1090,7 @@ subroutine dyn_core(npx, npy, npz, ng, sphum, nq, bdt, k_split, n_split, zvir, c
dudt_rf = ( u(is:ie,js:je,:) - dudt_rf)/dt
dvdt_rf = ( v(is:ie,js:je,:) - dvdt_rf)/dt
if (.not. hydrostatic) dwdt_rf = ( w(is:ie,js:je,:) - dwdt_rf)/dt
call timing_off('RAY_FAST')
endif

! *** Inline Beljaars turbulent-orographic-form-drag here
Expand Down Expand Up @@ -1195,20 +1236,11 @@ subroutine dyn_core(npx, npy, npz, ng, sphum, nq, bdt, k_split, n_split, zvir, c
!-----------------------------------------------------
enddo ! time split loop
!-----------------------------------------------------
if ( nq > 0 .and. .not. flagstruct%inline_q ) then
call timing_on('COMM_TOTAL')
call timing_on('COMM_TRACER')
call start_group_halo_update(i_pack(10), q, domain)
call timing_off('COMM_TRACER')
call timing_off('COMM_TOTAL')
endif


if ( flagstruct%fv_debug ) then
if(is_master()) write(*,*) 'End of n_split loop'
endif


if ( n_con/=0 .and. flagstruct%d_con > 1.e-5 ) then
nf_ke = min(3, flagstruct%nord+1)
call del2_cubed(heat_source, cnst_0p20*gridstruct%da_min, gridstruct, domain, npx, npy, npz, nf_ke, bd)
Expand Down
3 changes: 3 additions & 0 deletions model/fv_control.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -834,6 +834,9 @@ subroutine run_setup(Atm, dt_atmos, grids_on_this_pe, p_split)
write(*,199) 'Using p_split = ', p_split
endif

! Safety check on q_split and z_tracer
if ( q_split /= 0) z_tracer = .false.

if (Atm(n)%neststruct%nested) then
do i=1,n-1
if (Atm(i)%grid_number == parent_grid_num) then
Expand Down
106 changes: 43 additions & 63 deletions model/fv_dynamics.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,20 +245,15 @@ subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, &
type(fv_diag_type), intent(IN) :: idiag

! Local Arrays
#ifdef SINGLE_FV
! R8 Mass flux arrays: the "Flux Capacitor"
real(kind=8) :: mfxR8(bd%is:bd%ie+1, bd%js:bd%je, npz)
real(kind=8) :: mfyR8(bd%is:bd%ie , bd%js:bd%je+1, npz)
! R8 Courant number arrays
real(kind=8) :: cxR8(bd%is:bd%ie+1, bd%jsd:bd%jed, npz)
real(kind=8) :: cyR8(bd%isd:bd%ied ,bd%js:bd%je+1, npz)
#endif
! Local Mass flux arrays: the "Flux Capacitor"
real :: mfxL(bd%is:bd%ie+1, bd%js:bd%je, npz)
real :: mfyL(bd%is:bd%ie , bd%js:bd%je+1, npz)
! Local Courant number arrays
real :: cxL(bd%is:bd%ie+1, bd%jsd:bd%jed, npz)
real :: cyL(bd%isd:bd%ied ,bd%js:bd%je+1, npz)

integer :: imax_req
real :: cmax_z(npz)
! More Local arrays
real:: ws(bd%is:bd%ie,bd%js:bd%je)
real(kind=8):: te_2d(bd%is:bd%ie,bd%js:bd%je)
Expand All @@ -270,8 +265,7 @@ subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, &
real, dimension(bd%is:bd%ie):: cvm
real, allocatable :: dp1(:,:,:), dtdt_m(:,:,:), cappa(:,:,:)
real(kind=8), allocatable :: psx(:,:)
!real(kind=8), allocatable :: dpx(:,:)
real(kind=8), allocatable :: dpx(:,:,:) !needed for OpenMP
real(kind=8), allocatable :: dpx(:,:,:)
real:: akap, rdg, ph1, ph2, mdt, gam, amdt, u0
integer:: kord_tracer(ncnst)
integer :: i,j,k, n, iq, n_map, nq, nwat
Expand Down Expand Up @@ -604,14 +598,18 @@ subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, &
endif

!DryMassRoundoffControl
allocate(psx(isd:ied,jsd:jed),dpx(is:ie,js:je,npz))
psx(:,:) = 0.0
do j=js,je
do i=is,ie
psx(i,j) = pe(i,npz+1,j)
dpx(i,j,:) = 0.0
enddo
enddo
allocate(dpx(is:ie,js:je,npz))
dpx(:,:,:) = 0.0
if (hydrostatic) then
allocate(psx(isd:ied,jsd:jed))
psx(:,:) = 0.0
do j=js,je
do i=is,ie
psx(i,j) = pe(i,npz+1,j)
enddo
enddo
endif

call timing_on('FV_DYN_LOOP')
do n_map=1, k_split ! first level of time-split
call timing_on('COMM_TOTAL')
Expand Down Expand Up @@ -645,30 +643,12 @@ subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, &
call dyn_core(npx, npy, npz, ng, sphum, nq, mdt, k_split, n_split, zvir, cp_air, akap, cappa, grav, hydrostatic, &
u, v, w, delz, pt, q, delp, pe, pk, phis, varflt, ws, omga, ptop, pfull, ua, va, &
dudt_rf, dvdt_rf, dwdt_rf, uc, vc, &
#ifdef SINGLE_FV
mfxR8, mfyR8, cxR8, cyR8, &
#else
mfxL, mfyL, cxL, cyL, &
#endif
mfxR8, mfyR8, cxR8, cyR8, cmax_z, imax_req, q_split, &
pkz, peln, q_con, ak, bk, dpx, ks, &
gridstruct, flagstruct, neststruct, idiag, bd, &
domain, n_map==1, i_pack, last_step, diss_est,time_total)
call timing_off('DYN_CORE')

!MassFluxRoundoffControl
#ifdef SINGLE_FV
mfxL=mfxR8
mfyL=mfyR8
cxL= cxR8
cyL= cyR8
#endif
! if ( flagstruct%range_warn ) then
! call range_check('CX_dyn', cxL(is:ie,js:je,:)/real(n_split), is, ie, js, je, 0, npz, gridstruct%agrid, &
! -0.5, 0.5, bad_range)
! call range_check('CY_dyn', cyL(is:ie,js:je,:)/real(n_split), is, ie, js, je, 0, npz, gridstruct%agrid, &
! -0.5, 0.5, bad_range)
! endif

if ( flagstruct%fv_debug ) then
if(is_master()) write(6,*) ''
if(is_master()) write(6,*) '------------------------------------'
Expand All @@ -688,24 +668,25 @@ subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, &
!DryMassRoundoffControl
if(last_step) then
if (hydrostatic) then
do k = 2, npz
dpx(:,:,1) = dpx(:,:,1) + dpx(:,:,k)
enddo
do j=js,je
do i=is,ie
psx(i,j) = psx(i,j) + dpx(i,j,1)
enddo
enddo
call timing_on('COMM_TOTAL')
call mpp_update_domains(psx, domain)
call timing_off('COMM_TOTAL')
do j=js-1,je+1
do i=is-1,ie+1
pe(i,npz+1,j) = psx(i,j)
enddo
enddo
do k = 2, npz
dpx(:,:,1) = dpx(:,:,1) + dpx(:,:,k)
enddo
do j=js,je
do i=is,ie
psx(i,j) = psx(i,j) + dpx(i,j,1)
enddo
enddo
call timing_on('COMM_TOTAL')
call mpp_update_domains(psx, domain)
call timing_off('COMM_TOTAL')
do j=js-1,je+1
do i=is-1,ie+1
pe(i,npz+1,j) = psx(i,j)
enddo
enddo
deallocate(psx)
endif
deallocate(psx,dpx)
deallocate(dpx)
end if

#ifdef SW_DYNAMICS
Expand All @@ -723,18 +704,18 @@ subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, &
call timing_on('tracer_2d')
!!! CLEANUP: merge these two calls?
if (gridstruct%nested) then
call tracer_2d_nested(q, dp1, mfxL, mfyL, cxL, cyL, gridstruct, bd, domain, npx, npy, npz, nq, &
call tracer_2d_nested(q, dp1, mfxR8, mfyR8, cxR8, cyR8, cmax_z, imax_req, gridstruct, bd, domain, npx, npy, npz, nq, &
flagstruct%hord_tr, q_split, mdt, idiag%id_divg, i_pack(10), &
flagstruct%nord_tr, flagstruct%trdm2, &
k_split, neststruct, parent_grid, flagstruct%lim_fac)
else
if ( flagstruct%z_tracer .and. q_split == 0 ) then
call tracer_2d_1L(q, dp1, mfxL, mfyL, cxL, cyL, gridstruct, bd, domain, npx, npy, npz, nq, &
call tracer_2d_1L(q, dp1, mfxR8, mfyR8, cxR8, cyR8, cmax_z, imax_req, gridstruct, bd, domain, npx, npy, npz, nq, &
flagstruct%hord_tr, mdt, idiag%id_divg, i_pack(10), &
flagstruct%nord_tr, flagstruct%trdm2, flagstruct%lim_fac)
else
call tracer_2d(q, dp1, mfxL, mfyL, cxL, cyL, gridstruct, bd, domain, npx, npy, npz, nq, &
flagstruct%hord_tr, q_split, mdt, idiag%id_divg, i_pack(10), &
call tracer_2d(q, dp1, mfxR8, mfyR8, cxR8, cyR8, cmax_z, imax_req, gridstruct, bd, domain, npx, npy, npz, nq, &
flagstruct%hord_tr, flagstruct%n_sponge, q_split, mdt, idiag%id_divg, i_pack(10), &
flagstruct%nord_tr, flagstruct%trdm2, flagstruct%lim_fac)
endif
endif
Expand Down Expand Up @@ -793,7 +774,6 @@ subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, &
flagstruct%adiabatic, do_adiabatic_init, &
flagstruct%remap_option, flagstruct%gmao_remap, &
flagstruct%gmao_top_bc, flagstruct%gmao_bot_bc)
!!! mfx=mfxL, mfy=mfyL, cx=cxL, cy=cyL)

#ifdef AVEC_TIMERS
call avec_timer_stop(6)
Expand All @@ -808,10 +788,10 @@ subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, &
#endif

! Accumulate the total Mass flux and Courant numbers for export
mfx = mfx + mfxL
mfy = mfy + mfyL
cx = cx + cxL
cy = cy + cyL
mfx = mfx + mfxR8
mfy = mfy + mfyR8
cx = cx + cxR8
cy = cy + cyR8

if ( flagstruct%fv_debug ) then
if(is_master()) write(6,*) ''
Expand Down
Loading
Loading