Wrapper for communication resources

[devreal]

Non-contiguous Kokkos views may have to be either serialized (into a temporary buffer) or an MPI datatype has to be created and committed. Currently, these resources are managed on the fly, which may come at a cost. In #50 (comment) I suggested a wrapper that holds whatever kokkos deems necessary to perform the communication of a view.

// query whether we need a temporary buffer or a derived datatype
std::variant<size_t, KokkosComm::Type> opt = KokkosComm::query(view);
if (std::holds_alternative<size_t>(opt)) {
  // provide a temporary buffer for packing, takes ownership until completion
  ensure_size(std::get<size_t>(opt), my_properly_sized_buf);
  KokkosComm::isend(view, comm, ..., my_properly_sized_buf);
} else {
  // provides a datatype we have cached (could be built-in for contiguous views)
  KokkosComm::isend(view, comm, ..., std::get<KokkosComm::Type>(opt));
}

// no cached information for another_view, let kokkos figure it out
KokkosComm::isend(another_view, comm, ...,);

A slightly refined version of this could be:

// query whether we need a temporary buffer or a derived datatype
KokkosComm::Resource res = KokkosComm::query_resources(view);
KokkosComm::isend(view, comm, ..., res); // use datatype / buffer allocated above
MPI_Datatype dt = MPI_DATAYPE_NULL; 
if (rs.has_datatype()) res.datatype(); // get raw handle to datatype

// no cached information for another_view, let kokkos figure it out
KokkosComm::isend(another_view, comm, ...);

Alternatively, KokkosComm could cache the datatype in the session object that will be introduced later. That will not help with the buffer problem though.

A downside of the Resource handle is that it is not per-se thread-safe if a buffer is used instead of a datatype, unless buffers are allocated on a per-thread basis. This adds some complexity, either to the application or to the implementation.

[janciesko]

@nicoleavans, you might find this discussion interesting.

[cwpearson]

I think this is conceptually analogous to what I was working towards here: #64. I'm quite unsatisfied with the interface I came up with for the Ctx (aka Resource) object I wrote though

A problem I ran in to is that it felt like point-to-point vs various collectives really each wanted Resource to have a slightly different interface, since functions that take multiple buffers have intertwined semantics that need some coordination across the non-contiguous strategy used for those buffers.

The interface I propose currently supports splitting a message into multiple messages (I was thinking ahead to partitioned communication)

[cwpearson]

I have a revised attempt, no PR yet, but the conclusion is something like this (from 2d halo perf test). It's kind of analogous to the Group interface in NCCL (@dssgabriel)

diff --git a/perf_tests/test_2dhalo.cpp b/perf_tests/test_2dhalo.cpp
index 8b34381..638f267 100644
--- a/perf_tests/test_2dhalo.cpp
+++ b/perf_tests/test_2dhalo.cpp
@@ -46,12 +46,12 @@ void send_recv(benchmark::State &, MPI_Comm comm, const Space &space, int nx, in
   auto ym1_s = Kokkos::subview(v, make_pair(1, nx + 1), 1, Kokkos::ALL);
   auto ym1_r = Kokkos::subview(v, make_pair(1, nx + 1), 0, Kokkos::ALL);
 
-  std::vector<KokkosComm::Req> reqs;
-  // std::cerr << get_rank(rx, ry) << " -> " << get_rank(xp1, ry) << "\n";
-  reqs.push_back(KokkosComm::isend(space, xp1_s, get_rank(xp1, ry), 0, comm));
-  reqs.push_back(KokkosComm::isend(space, xm1_s, get_rank(xm1, ry), 1, comm));
-  reqs.push_back(KokkosComm::isend(space, yp1_s, get_rank(rx, yp1), 2, comm));
-  reqs.push_back(KokkosComm::isend(space, ym1_s, get_rank(rx, ym1), 3, comm));
+  KokkosComm::Handle<Space> h = KokkosComm::plan(space, comm, [=](KokkosComm::Handle<Space> &handle) {
+    KokkosComm::isend(handle, xp1_s, get_rank(xp1, ry), 0);
+    KokkosComm::isend(handle, xm1_s, get_rank(xm1, ry), 1);
+    KokkosComm::isend(handle, yp1_s, get_rank(rx, yp1), 2);
+    KokkosComm::isend(handle, ym1_s, get_rank(rx, ym1), 3);
+  });
 
   KokkosComm::recv(space, xm1_r, get_rank(xm1, ry), 0, comm);
   KokkosComm::recv(space, xp1_r, get_rank(xp1, ry), 1, comm);
@@ -59,9 +59,7 @@ void send_recv(benchmark::State &, MPI_Comm comm, const Space &space, int nx, in
   KokkosComm::recv(space, yp1_r, get_rank(rx, yp1), 3, comm);
 
   // wait for comm
-  for (KokkosComm::Req &req : reqs) {
-    req.wait();
-  }
+  h.wait();
 }

I think the names of the types are a work in progress, but this fixes the problem with the current implementation where isend has no choice but to fence before each call. On CPU the fence is a no-op, which is okay, but on GPU it's not.

With this construction, the communicator and space are associated with the handle, rather than each operation. Those isend / irecv calls register various operations they want with the handle, and then the handle can minimize the number of required fences. Similarly, this handle has enough flexibility internally to tolerate a variety of non-contiguous data handling strategies.

[cwpearson]

as seen in #100