Skip to content

Scheduled Halo Exchange #980

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
philip-paul-mueller wants to merge 29 commits into
base: main
Choose a base branch
from
Open

Scheduled Halo Exchange #980

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
26d685b
Modified versions.
philip-paul-mueller Dec 18, 2025
6518ce9
Made some addaptions towards the asynchronous exchange.
philip-paul-mueller Dec 18, 2025
ec7fca2
More uniformity.
philip-paul-mueller Dec 18, 2025
1f5e9e6
Updated ghex version.
philip-paul-mueller Dec 18, 2025
e69cb82
Fixed at least that issue.
philip-paul-mueller Dec 18, 2025
f60a1f8
Made the components aware of async stuff.
philip-paul-mueller Dec 18, 2025
e11da41
Fixed some stray `stream` argument.
philip-paul-mueller Dec 19, 2025
6636cea
Updated the annotations and the meaning of `stream`.
philip-paul-mueller Dec 19, 2025
383f959
Realized that the strams are disabled.
philip-paul-mueller Dec 19, 2025
41322f7
Let's see if that help, but it is strange that it takes longer now.
philip-paul-mueller Dec 19, 2025
ae6db39
Updated ghex version.
philip-paul-mueller Dec 19, 2025
815fc46
This should fix the argument names, but I do not understant why it is…
philip-paul-mueller Dec 19, 2025
296c7b1
Renambled named arguments.
philip-paul-mueller Dec 19, 2025
310a5e5
There it is not \(yet\) possible to use the name.
philip-paul-mueller Dec 19, 2025
9a8975f
Made it mandatory to pass the stream.
philip-paul-mueller Dec 23, 2025
a985a59
Undo something.
philip-paul-mueller Dec 23, 2025
e04e057
Updated GHEX.
philip-paul-mueller Dec 23, 2025
7fc5154
Forgot one.
philip-paul-mueller Dec 23, 2025
9a68a54
Let's hope that is enough.
philip-paul-mueller Dec 23, 2025
0d3a400
Added warnings.
philip-paul-mueller Dec 24, 2025
6f4c6da
Updated GHEX.
philip-paul-mueller Dec 24, 2025
aae7b06
Starnge things are going on.
philip-paul-mueller Dec 24, 2025
d6dbc8c
The problem was me installing GHEX wrong.
philip-paul-mueller Dec 24, 2025
c8d58f8
The no streaming case is no longer `None` but has a named constant.
philip-paul-mueller Jan 1, 2026
89a94aa
Forgot to add a warning.
philip-paul-mueller Jan 1, 2026
0b8307b
I am not fully like it, but for this it seems appropriate.
philip-paul-mueller Jan 1, 2026
0e2fbde
Forgot that.
Jan 2, 2026
0694826
Merge remote-tracking branch 'gitlab/main' into phimuell__async_mpi_test
Jan 2, 2026
a570aee
Removed useless return.
Jan 2, 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
16 changes: 13 additions & 3 deletions model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,9 @@ def run(
log.debug("advection run - start")

log.debug("communication of prep_adv cell field: mass_flx_ic - start")
self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic)
self._exchange.exchange_and_wait(
dims.CellDim, prep_adv.mass_flx_ic, stream=decomposition.DefaultStream
)
log.debug("communication of prep_adv cell field: mass_flx_ic - end")

log.debug("running stencil copy_cell_kdim_field - start")
Expand Down Expand Up @@ -271,7 +273,11 @@ def run(
log.debug("advection run - start")

log.debug("communication of prep_adv cell field: mass_flx_ic - start")
self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic)
self._exchange.exchange_and_wait(
dims.CellDim,
prep_adv.mass_flx_ic,
stream=decomposition.DefaultStream,
)
log.debug("communication of prep_adv cell field: mass_flx_ic - end")

# reintegrate density for conservation of mass
Expand Down Expand Up @@ -364,7 +370,11 @@ def run(

# exchange updated tracer values, originally happens only if iforcing /= inwp
log.debug("communication of advection cell field: p_tracer_new - start")
self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new)
Copy link
Contributor

@muellch muellch Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have to wait here?
This is not the end of the dycore, just the end of the corrector.

There are more dycore stencils after this, and the wait could happen here:

icon4py/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py

Line 1108 in a76e26b

This makes the code uglier though.

Copy link
Collaborator Author

@philip-paul-mueller philip-paul-mueller Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure I would say this is something for another PR.
The main question is if it should happen before or after this PR is merged.

Copy link
Collaborator Author

@philip-paul-mueller philip-paul-mueller Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for looking into it.

self._exchange.exchange_and_wait(
dims.CellDim,
p_tracer_new,
stream=decomposition.DefaultStream,
)
log.debug("communication of advection cell field: p_tracer_new - end")

# finalize step
Expand Down
4 changes: 3 additions & 1 deletion model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,9 @@ def apply_flux_limiter(
)

log.debug("communication of advection cell field: r_m - start")
self._exchange.exchange_and_wait(dims.CellDim, self._r_m)
self._exchange.exchange_and_wait(
dims.CellDim, self._r_m, stream=decomposition.DefaultStream
)
log.debug("communication of advection cell field: r_m - end")

# limit outward fluxes
Expand Down
25 changes: 21 additions & 4 deletions model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -386,7 +386,7 @@ def __init__(
self._cell_params = cell_params

self.halo_exchange_wait = decomposition.create_halo_exchange_wait(
self._exchange
self._exchange,
) # wait on a communication handle
self.rd_o_cvd: float = constants.GAS_CONSTANT_DRY_AIR / (
constants.CPD - constants.GAS_CONSTANT_DRY_AIR
Expand Down Expand Up @@ -736,6 +736,7 @@ def _sync_cell_fields(self, prognostic_state):
prognostic_state.w,
prognostic_state.theta_v,
prognostic_state.exner,
stream=decomposition.DefaultStream,
)
log.debug("communication of prognostic cell fields: theta, w, exner - done")

Expand Down Expand Up @@ -772,12 +773,14 @@ def _do_diffusion_step(
log.debug("rbf interpolation 1: end")

# 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert
# TODO(phimuell, muellch): Is asynchronous mode okay here.
log.debug("communication rbf extrapolation of vn - start")
self._exchange(
self.u_vert,
self.v_vert,
dim=dims.VertexDim,
wait=True,
stream=decomposition.DefaultStream,
)
log.debug("communication rbf extrapolation of vn - end")

Expand Down Expand Up @@ -817,7 +820,13 @@ def _do_diffusion_step(
# TODO(halungge): move this up and do asynchronous exchange
if self.config.type_vn_diffu > 1:
log.debug("communication rbf extrapolation of z_nable2_e - start")
self._exchange(self.z_nabla2_e, dim=dims.EdgeDim, wait=True)
# TODO(phimuell, muellch): Is asynchronous mode okay here.
self._exchange(
self.z_nabla2_e,
dim=dims.EdgeDim,
wait=True,
stream=decomposition.DefaultStream,
)
log.debug("communication rbf extrapolation of z_nable2_e - end")

log.debug("2nd rbf interpolation: start")
Expand All @@ -827,12 +836,14 @@ def _do_diffusion_step(
log.debug("2nd rbf interpolation: end")

# 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields)
# TODO(phimuell, muellch): Is asynchronous mode okay here.
log.debug("communication rbf extrapolation of z_nable2_e - start")
self._exchange(
self.u_vert,
self.v_vert,
dim=dims.VertexDim,
wait=True,
stream=decomposition.DefaultStream,
)
log.debug("communication rbf extrapolation of z_nable2_e - end")

Expand All @@ -848,7 +859,12 @@ def _do_diffusion_step(
log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end")

log.debug("communication of prognistic.vn : start")
handle_edge_comm = self._exchange(prognostic_state.vn, dim=dims.EdgeDim, wait=False)
handle_edge_comm = self._exchange(
prognostic_state.vn,
dim=dims.EdgeDim,
wait=False,
stream=decomposition.DefaultStream,
)

log.debug(
"running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start"
Expand Down Expand Up @@ -894,7 +910,8 @@ def _do_diffusion_step(
log.debug("running stencil 13 to 16 apply_diffusion_to_theta_and_exner: end")

self.halo_exchange_wait(
handle_edge_comm
handle_edge_comm,
stream=decomposition.DefaultStream,
) # need to do this here, since we currently only use 1 communication object.
log.debug("communication of prognogistic.vn - end")

Expand Down
23 changes: 19 additions & 4 deletions model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1190,7 +1190,10 @@ def run_predictor_step(

log.debug("exchanging prognostic field 'vn' and local field 'rho_at_edges_on_model_levels'")
self._exchange.exchange_and_wait(
dims.EdgeDim, prognostic_states.next.vn, z_fields.rho_at_edges_on_model_levels
dims.EdgeDim,
prognostic_states.next.vn,
z_fields.rho_at_edges_on_model_levels,
stream=decomposition.DefaultStream,
)

self._compute_horizontal_velocity_quantities_and_fluxes(
Expand Down Expand Up @@ -1262,11 +1265,18 @@ def run_predictor_step(
"exchanging prognostic field 'w' and local field 'dwdz_at_cells_on_model_levels'"
)
self._exchange.exchange_and_wait(
dims.CellDim, prognostic_states.next.w, z_fields.dwdz_at_cells_on_model_levels
dims.CellDim,
prognostic_states.next.w,
z_fields.dwdz_at_cells_on_model_levels,
stream=decomposition.DefaultStream,
)
else:
log.debug("exchanging prognostic field 'w'")
self._exchange.exchange_and_wait(dims.CellDim, prognostic_states.next.w)
self._exchange.exchange_and_wait(
dims.CellDim,
prognostic_states.next.w,
stream=decomposition.DefaultStream,
)

def run_corrector_step(
self,
Expand Down Expand Up @@ -1361,7 +1371,11 @@ def run_corrector_step(
)

log.debug("exchanging prognostic field 'vn'")
self._exchange.exchange_and_wait(dims.EdgeDim, (prognostic_states.next.vn))
self._exchange.exchange_and_wait(
dims.EdgeDim,
prognostic_states.next.vn,
stream=decomposition.DefaultStream,
)

self._compute_averaged_vn_and_fluxes_and_prepare_tracer_advection(
spatially_averaged_vn=self.z_vn_avg,
Expand Down Expand Up @@ -1433,4 +1447,5 @@ def run_corrector_step(
prognostic_states.next.rho,
prognostic_states.next.exner,
prognostic_states.next.w,
stream=decomposition.DefaultStream,
)
Loading