Added `schedule_exchange()` and `schedule_wait()` to `communication_object` by philip-paul-mueller · Pull Request #190 · ghex-org/GHEX

philip-paul-mueller · 2025-12-18T07:08:38Z

This PR adds the the schedule_exchange() and schedule_wait() function to the communication object.
They behave similar to the regular exchange() and wait() but they accepts an additional CUDA stream as argument.

schedule_exchange() will wait with packing until all work that has been scheduled on the passed stream has finished, which removes the need for an external synchronization.
It is important to note, that the function will only return after all data have been send, which is the same behaviour than exchange().

schedule_wait() is similar to wait(), it will launch unpacking but will make synchronize it with the passed stream.
This means that every work that is submitted to the stream after the function has returned, will not start until all unpacking has been completed.

The PR also extends the Python bindings.
The bindings are able to interpret the following Python objects:

None is interpreted as default stream, i..e nullptr.
If the object has a __cuda_stream__ method it is assumed to follow Nvidia's stream protocol.
If the object has a .ptr attribute it is assumed that it follows CuPy's Stream implementation.

Note:
For CPU memory schedule_exchange() and schedule_wait() behave the exact same way as exchange() and wait(), i.e. exchange() starts to pack immediately and wait() only returns after the unpacking has finished.
If CPU and GPU memory is exchanged in the same transaction then the behaviour is a mix of both, the CPU parts are immediately packed but the packing of the GPU memory synchronizes with the stream.
The same holds for schedule_wait(), which will only return if the CPU memory has been unpacked, but will return after the unpacking of the GPU memory has been initiated.
What happens exactly also depends on if the CPU memory is processed before the GPU memory.
Thus it is safe to mix the two but it is not recommended.

NOTE:

This PR replaces PR#186, with a refactored solution.
The main important changes to this function are in the following files:
- include/ghex/communication_object.hpp: Adding of the schedule_*() functions.
- include/ghex/device/cuda/stream.hpp: Adding of a (very simple) event pool.
- bindings/python/src/_pyghex/unstructured/communication_object.cpp: Updating the bindings.
This PR depends on PR#189.
This PR depends on PR#191, which adds formatting.

TODO:

This PR still has the wrong formatting applied to it.
Test on ICON4Py production setup.
Ask Fabian if we should add CuPy as a dependency for the binding tests such that we can also test the GPU tests.

…t compiles though.

…here is another error.

msimberg · 2026-01-07T11:22:17Z

include/ghex/communication_object.hpp

+     *
+     * TODO: Should the handle expose this function?
+     */
+    void complete_schedule_exchange()


What do you think about making this private and if a caller needs to synchronize after a schedule_wait, just make them call wait instead?

schedule_exchange(stream); schedule_wait(stream); // now I really want to synchronize wait();

IMO this would simplify the API in that wait can always be called as a wait-for-everything-to-finish regardless of how the exchange or wait was done previously. wait may of course need to call complete_schedule_exchange internally (or parts of it) for this approach to work.

We should consider it.

I have though about it again.

Regarding your example, while I can understand why it should work (in fact I also think it must work), I do not think that it is the right (what ever that means) way to achieve it.
The scheduled functions essentially introduce a kind of stream semantic into GHEX.
So, I would argue that to check if the transmission has finished one should synchronize with the stream passed to schedule_exchange().
The only thing that GHEX must do is, not start a new exchange before the old has finished or delete something that is still in use by a transfer that is still ongoing, which is what complete_schedule_exchange() does.
So I think that you are right when you say that it should become private as it should probably never be called directly by a user.
Instead all exchange*() functions call it to make sure that the exchange has happened and it is safe to start a new one.
Because wait() deallocate memory it must also call it to make sure that it does not delete something that is still in use.
As a side effect your example code will work and do what you want, i.e. the full synchronization.
But, as I outlined above, in my opinion, it is not the right way of doing it, but it must work.

Does this make sense?

I think that it should be private but currently the test needs them and currently I have no good idea to solve it.

msimberg · 2026-01-08T14:13:44Z

include/ghex/communication_object.hpp

+#ifdef GHEX_CUDACC
+        assert(has_scheduled_exchange());
+#endif


This doesn't belong here, or the logic is wrong elsewhere. clear() is called by complete_scheduled_exchange after the event has been reset, meaning this always fails. Should this be assert(!has_scheduled_exchange())?

You are right, it should be assert(!...), I am just wondering why the CI "passes".
Probably because asserts are disabled?

msimberg · 2026-01-08T14:14:01Z

include/ghex/communication_object.hpp

+     * exchange. A user will never have to call it directly. If there was no such
+     * exchange or GPU support was disabled, the function does nothing.
+     *
+     * \note  This should be a private function, but the tests need them.


Why do the tests need this function?

Some test calls it (https://github.com/ghex-org/GHEX/pull/190/changes#diff-c9ca950ae3be0eda232ecea6c5e5224c87bdee091f1e3a7c9efa081badda3fc2R364) but I think that this is more like a bad test.

…ecome private.

boeschf

First of all, thank you for the many improvements! This is a much needed addition.

There are a few nitpicks on which I commented. Feel free to postpone them/ignore them.

In general, there are a few missing pieces still, which could be addressed in the future:

only the unstructured part is exposed to Python, but not the structured
the bulk_communication_object could perhaps also benefit from the enriched interface

Overall, the communication object becomes quite an unwieldy beast, and we should start thinking about moving to a true async runtime.

boeschf · 2026-01-13T08:59:02Z

include/ghex/device/cuda/event.hpp

+        if (!m_moved) { GHEX_CHECK_CUDA_RESULT_NO_THROW(cudaEventDestroy(m_event)) }
+    }
+
+    operator bool() const noexcept { return m_moved; }


I agree that the semantics are backwards here - this should be changed everywhere, but perhaps in another PR.

boeschf · 2026-03-13T15:26:01Z

include/ghex/communication_object.hpp

+        // Make sure that communication has finished and we can deallocate
+        // the buffers. Maybe the call to `clear()` is too much here and
+        // we should only wait.
+        complete_schedule_exchange();


comments seem not to match the code

The part about clear() was about what happens inside the complete_schedule_exchange() and was more like a note to myself.

boeschf · 2026-03-13T15:35:53Z

include/ghex/communication_object.hpp

+    template<bool UseAsyncStream, typename... StreamType>
+    void pack_and_send_impl(StreamType&&... sync_streams)


It seems the template parameter pack is only needed to allow for either zero or exactly one argument - this is not really nice, and entails the whole array expansion below. I would simplify this (enable_if for example), but I am ok with this considering this is an implementation detail.

I did it to allow multiple streams, but realized that it was a bit hard to do so I abandoned this path.

I have now separated the two implementation, it is a little bit code dublication but not much.

philip-paul-mueller · 2026-03-16T07:42:04Z

@boeschf Thanks for your comments.
I think we should discuss the more general ones tomorrow in person.

boeschf

thanks, looks good to me

msimberg · 2026-03-17T12:30:42Z

I will merge this now to prepare a new release. Further cleanups are still useful, but will be part of #195 and other PRs for the release after this one.

This PR introduces the [scheduled exchange feature](ghex-org/GHEX#190) from GHEX into ICON4Py. These exchange allows to call the exchange function before all work has been completed, i.e. the exchange will wait until the previous work is done. A similar feature is the "scheduled wait", that allows to initiate the receive without the need to wait on its completion. In addition to this the function also renamed the functions related to halo exchange: - `exchange()` was renamed to `start()`. - `wait()` was renamed to `finish()` (that might now return before the transfer has fully concluded). - `exchange_and_wait()` was renamed to `exchange()`. All of these functions now accepts the an argument called `stream`, which defaults to `DEFAULT_STREAM`. It is indicate how synchronization with the stream should be performed. In case of `start()` it means that the actual exchange should not start until all work previously submitted to `stream` has finished. For `finish()` it means that further work, submitted to `stream`, should not start until the exchange has ended. For `finish()` it is also possible to specify `BLOCK`, which means that `finish()` waits until the transfer has fully finished. The orchestrator was not updated, but the change were made in such a way that it continues to work in diffusion, although using the original, blocking behaviour. --------- Co-authored-by: Philip Mueller <[email protected]> Co-authored-by: Hannes Vogt <[email protected]> Co-authored-by: Copilot Autofix powered by AI <[email protected]> Co-authored-by: Mikael Simberg <[email protected]> Co-authored-by: Hannes Vogt <[email protected]>

@philip-paul-mueller

Adds gtfn_gpu backend to the distributed CI pipeline. dace_gpu is still left out because compilation takes too long. The base image is upgraded because it's possible, but not strictly necessary. The CPU-only version of the pipeline needed 25.04 (24.04 and 25.10 did not work for various reasons). However, since OpenMPI and libfabric are now built manually in the container the base image version is less of a constraint. 24.04 doesn't have matching GCC/CUDA versions and 26.04 doesn't exist yet, but the pipeline should eventually use 26.04. OpenMPI and libfabric are built manually for slingshot support because getting the ubuntu repository packages to work with GPU support did not seem possible/easy. The installation is based on https://github.com/eth-cscs/cray-network-stack. GHEX needs an upgrade, because there's a bug in how strides are calculated for GPU buffers. @philip-paul-mueller has already fixed this in ghex-org/GHEX#190 but we should wait for that to be merged (and probably test in icon-exclaim first). This also fixes a few cupy/numpy incompatibilities. `revert_repeated_index_to_invalid` was updated to only deal with numpy for now as the connectivities are always numpy arrays. `test_halo_exchange_for_sparse_field` is marked `embedded_only`. The non-MPI test was already marked embedded-only. This does not try to unify the default and distributed CI pipeline definitions. That should, however, be done done sooner or later as well. --------- Co-authored-by: Jacopo Canton <[email protected]> Co-authored-by: Nicoletta Farabullini <[email protected]>

msimberg and others added 30 commits November 3, 2025 12:48

Don't wait for streams to finish unpacking

1add6ed

Add dependency on default stream before starting packing

6d89616

hacky async mpi changes

714f7b3

Made some notes for me and a plan forward.

06dec72

It seems a bit overkill, but let's see if it works, or even compile.

9ef0f07

Removed the stream member from teh comunication object, not sure if i…

336117a

…t compiles though.

Run the code formater.

e1f918f

Fixed smaller things.

d7ff9a3

I have to ask somebody about that.

6df7291

Update the description.

098a3f3

First version of the event pool.

3ad6b17

Forgot to update something.

e4f35d7

Updated some things.

7b05329

Applied some changes after discussing them with Mikael, but I think t…

d2066a6

…here is another error.

Fixed some bugs, but I am not sure if it compiles and is correct.

633b17c

Small update.

1a65634

The python interface now accepts streams.

dfd7065

Applied the formatter.

e95665a

Fixed some issues, this should probably be enough.

30ed9dc

Added streams to the code.

315c647

Made it such that one can switch between default and non default stream.

2efacb1

Made the test 'GPU aware' not realy some pieces are missing.

c07f7f5

This should make it work on GPU.

791725f

Forgot to update the 'schedule_wait' call.

a87178d

Fixed a bug in strides computation, in default case.

63c05f9

Fixed another issue.

9b0f17a

Modified the checking a bit.

f1e72d0

This should be more GPU aware.

9a0d35f

Made more verbose error messages.

c04ab31

Why does this bug always happens to me?

0612e95

philip-paul-mueller added 4 commits January 6, 2026 07:04

Forgot to apply the formating.

78486e5

Forgot some stuff.

82fbb3b

Why do I forgot them all the time.

9cc59bd

I think I should start compiling it locally.

e4277d5

msimberg reviewed Jan 7, 2026

View reviewed changes

philip-paul-mueller added 5 commits January 8, 2026 08:08

Applied Mikaels comments.

68751a5

Forgot to update the descripton.

d50db32

This function should be public.

51e20d2

Forgot to update the bindings.

34141c4

Make compelete_schedule_exchange() again a private function.

7825d17

I think that it should be private but currently the test needs them and currently I have no good idea to solve it.

msimberg reviewed Jan 8, 2026

View reviewed changes

Modified teh code such that complete_scheduled_exchange() can now b…

80c0650

…ecome private.

philip-paul-mueller marked this pull request as ready for review January 23, 2026 09:10

msimberg mentioned this pull request Feb 4, 2026

Add gtfn_gpu to distributed CI pipeline C2SM/icon4py#1012

Merged

philip-paul-mueller mentioned this pull request Feb 6, 2026

Scheduled Halo Exchange C2SM/icon4py#980

Merged

Some small modification of comments.

a1de2a6

boeschf approved these changes Mar 13, 2026

View reviewed changes

Applied some comments.

d8d69ba

philip-paul-mueller requested a review from boeschf March 16, 2026 07:42

philip-paul-mueller added 2 commits March 16, 2026 08:53

Forgot to remove old code.

4e9a7cd

What are you doing CI.

9c09f3c

boeschf approved these changes Mar 16, 2026

View reviewed changes

msimberg merged commit 871b380 into ghex-org:master Mar 17, 2026
11 of 12 checks passed

msimberg mentioned this pull request Mar 17, 2026

Bump version to 0.5.0 #197

Merged

msimberg mentioned this pull request Mar 20, 2026

Invert meaning of operator bool for CUDA stream and event #199

Open

		template<bool UseAsyncStream, typename... StreamType>
		void pack_and_send_impl(StreamType&&... sync_streams)

Conversation

philip-paul-mueller commented Dec 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

philip-paul-mueller Jan 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

philip-paul-mueller Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

boeschf left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

philip-paul-mueller Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

philip-paul-mueller commented Mar 16, 2026

Uh oh!

boeschf left a comment

Choose a reason for hiding this comment

Uh oh!

msimberg commented Mar 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

philip-paul-mueller commented Dec 18, 2025 •

edited

Loading

philip-paul-mueller Jan 7, 2026 •

edited

Loading

philip-paul-mueller Jan 9, 2026 •

edited

Loading

philip-paul-mueller Mar 16, 2026 •

edited

Loading