feat(mpi): add host staging for ibroadcast

Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>

Author: dssgabriel

src/KokkosComm/mpi/broadcast.hpp

@@ -12,6 +12,7 @@
 #include "mpi_space.hpp"
 #include "req.hpp"
 
+#include <KokkosComm/impl/host_staging.hpp>
 #include "impl/error_handling.hpp"
 
 namespace KokkosComm {
@@ -23,12 +24,20 @@ auto ibroadcast(const ExecSpace& space, View& v, int root, MPI_Comm comm) -> Req
   Kokkos::Tools::pushRegion("KokkosComm::mpi::ibroadcast");
   fail_if(!is_contiguous(v), "KokkosComm::mpi::ibroadcast: unimplemented for non-contiguous views");
 
-  // Sync: Work in space may have been used to produce view data.
-  space.fence("fence before non-blocking broadcast");
-
   Req<MpiSpace> req;
-  MPI_Ibcast(data_handle(v), span(v), datatype<MpiSpace, T>, root, comm, &req.mpi_request());
+#if defined(KOKKOSCOMM_ENABLE_GPU_AWARE_MPI)
+  // Sync: Work in space may have been used to produce view data.
+  space.fence("fence before GPU-aware `MPI_Ibcast`");
+  MPI_Ibcast(data_handle(v), span(v), datatype<MpiSpace, T>(), root, comm, &req.mpi_request());
   req.extend_view_lifetime(v);
+#else
+  auto host_v = KokkosComm::Impl::stage_for(v);
+  // Sync: Ensure that `host_v` is done being copied on the host
+  space.fence("fence before non-blocking broadcast");
+  MPI_Ibcast(data_handle(host_v), span(host_v), datatype<MpiSpace, T>(), root, comm, &req.mpi_request());
+  // Implicitly extends lifetimes of `host_v` and `v` due to lambda capture
+  req.call_after_mpi_wait([=]() { KokkosComm::Impl::copy_back(space, v, host_v); });
+#endif
 
   Kokkos::Tools::popRegion();
   return req;
@@ -48,7 +57,7 @@ void broadcast(ExecSpace const& space, View const& v, int root, MPI_Comm comm) {
   auto host_v = KokkosComm::Impl::stage_for(v);
   // Sync: Ensure that `host_v` is done being copied on the host
   space.fence("fence before non-blocking broadcast");
-  MPI_Bcast(data_handle(host_v), span(host_v), datatype<MpiSpace, T>, root, comm);
+  MPI_Bcast(data_handle(host_v), span(host_v), datatype<MpiSpace, T>(), root, comm);
   KokkosComm::Impl::copy_back(space, v, host_v);
 #endif