Fencing and async communication

[cwpearson]

Option 1

req = isend(space, ...);
req.wait();

// or, without spaces
req = isend(...);
req.wait();

For the above, to ensure communication is actually done, req.wait() may in general need to fence space, which means Req needs to have an execution space member in it. Since Req is not a template struct (as currently proposed), that member execution space would need to be type-erased.

We could sidestep this by requiring that the user fence any space instances they explicitly use. This is kind of a bummer because now wait is not sufficient to finish the communication.

Option 2

req = isend(space, ...)
space.fence();
req.wait();
// or space.fence() here

// or, without space instances
req = isend();
Kokkos::fence();
req.wait();

Discussion

The challenege with option 1 is that it makes the following pattern hard to implement:

r1 = isend(space, ...)
r2 = isend(space, ...)

wait_all(r1, r2, ...); // it would be cool if this only called space.fence() once

For wait_all to know that space.fence only needs to be called once, wait_all needs to see the actual execution space instances held in the Req, but it can't because they've been type-erased. We'd have to use RTTI to tag the Req with the type of the held space, and then do some dynamic casting and call the space's operator== to issue only one fence for each space instance present across all requests.

Or Req could be templated on the Kokkos Execution space. Maybe this is fine - we could do something like a HostReq and a Req for Requests against host and device spaces.

[devreal]

Option 3**

Define the req.wait to be ordered within the execution space and require the application to fence if it needs to. That way, asynchronous data movements (e.g., the unpack of a message) on a CUDA stream for example don't have to be fenced if the next kernel is submitted on that same stream. This would be true for NCCL (if you want to use that as a backend) and may one day be true for MPI as well. In what situations would the user have to fence the output of the communication anyway? Isn't everything ordered inside an execution context? (this may be my ignorance of kokkos semantics, apologies)

[devreal]

Thinking about it: req.wait() seems utterly misplaced in this scheme. I would argue that it should be KokkosComm::wait (and KokkosComm::waitall/any/some if wanted). There are even ways to synchronize streams if the isend/irecv were started on one stream and the wait happens on another (using events).

In general, I'd much prefer free-standing functions over member functions in this whole API.

[cwpearson]

Very helpful, I think this may be the right way to proceed. My only complaint is that req.wait() may implicitly interact with space (right?) rather than explicitly, but I think it's still better than the two options I gave. This could be made a bit clearer if Req ended up templated on the execution space type - the user would understand there is some kind of interaction there. Also, this proposal could work well in a future world with stream-triggered communication where wait would not need to block the host thread.

In general, I'd much prefer free-standing functions over member functions in this whole API.

Yes - maybe KokkoComm::wait(req) would be better. I noticed you wrote KokkosComm::isend(space, ..., &req); is there a reason you prefer that to req = ...?

[devreal]

Yes - maybe KokkoComm::wait(req) would be better. I noticed you wrote KokkosComm::isend(space, ..., &req); is there a reason you prefer that to req = ...?

Ahh no, that was my C brain writing C MPI code :) returning the request is better.

[cwpearson]

If wait doesn't necessarily block the host, maybe it needs a different name to avoid confusion with MPI_Wait? wait_async for nonblocking, and wait for blocking?

[devreal]

It will block the execution space until the operation is complete. That may entail blocking the host thread. I think wait is ok but it needs to be documented.

[dutkalex]

Option 2 is very error-prone IMO, because the behavior of wait becomes counter-intuitive... In our in-house solution, we've had some success tackling this issue by decoupling the wrapping of the MPI_Request and the communication scheme. I would argue that the type erased solution might be a good default, but it leaves performance on the table in the case where all the relevant information is known at compile time, which would justify to have a different implementation available for those cases... I would further argue that separating the packing part from the sending part would enable the implementation of more efficient communication schemes.

[devreal]

I agree that there should be a low-level API or some way to provide resources to the high-level API to avoid allocation/deallocation on the critical path.

[cwpearson]

What's the scope of this proposed low-level API? Just any contiguous 1D view, and otherwise basically identical to existing MPI?

[dutkalex]

Yes. I think it should be easily extensible to multidimensional views as long as the layout is contiguous. Maybe other zero overhead extensions are possible but nothing comes to my mind right now

[cwpearson]

Nothing wrong with this very simple wrapper - feel free to make PRs especially for any functions you're particularly fond of 😄

[dutkalex]

I will convert #50 into a real PR then :)

[cwpearson]

I'm kind of inclined to template Req on the execution space. Most codes will only ever interact with two execution space types simultaneously, so where they currently only need one container of requests, they would need two and ONLY if they are using two different space types for communication at once