feat: design a generic memory space staging API

[dssgabriel]

This PR introduces a minimal API that allows us to perform host staging whenever the underlying communication backend does not support GPU-aware operations (currently, this can only occur with MPI):

• stage_for: creates a mirror view if the passed view type isn't host-accessible, otherwise returns the same view
• copy_back: deep copies the host mirror view back to the device view when it is not host accessible (required for receive operations), otherwise does nothing
  I am restricting this API to internal use for now, hence it is behind the Impl:: namespace.

This is a critical feature that is needed to properly fall back to pure CPU-based communications in order to support MPI routines that are not yet GPU-aware (e.g. non-blocking collectives in Open MPI).

Usage would look something like so (here in mpi::iallreduce):

template <KokkosView SView, KokkosView RView, KokkosExecutionSpace ExecSpace>
auto iallreduce(const ExecSpace &space, const SView sv, RView rv, MPI_Op op, MPI_Comm comm) -> Req<MpiSpace> {
  using ST = typename SView::non_const_value_type;
  using RT = typename RView::non_const_value_type;
  static_assert(std::is_same_v<ST, RT>, "KokkosComm::mpi::iallreduce: View value types must be identical");
  Kokkos::Tools::pushRegion("KokkosComm::mpi::iallreduce");

  fail_if(!is_contiguous(sv) || !is_contiguous(rv),
          "KokkosComm::mpi::iallreduce: unimplemented for non-contiguous views");

  // Sync: Work in space may have been used to produce view data.
  space.fence("fence before non-blocking all-gather");

  Req<MpiSpace> req;
  if constexpr (Impl::is_gpu_aware()) {  // NOTE: this API does not exist (yet)
    MPI_Iallreduce(data_handle(sv), data_handle(rv), span(sv), datatype<MpiSpace, ST>(), op, comm, &req.mpi_request());
  } else {
    auto host_staged_sv = KokkosComm::Impl::stage_for(space, sv);
    auto host_staged_rv = KokkosComm::Impl::stage_for(space, rv);
    space.fence();
    MPI_Iallreduce(data_handle(host_staged_sv), data_handle(host_staged_rv), span(host_staged_sv), datatype<MpiSpace, ST>(), op, comm, &req.mpi_request());
    req.call_after_mpi_wait([=]() {
      KokkosComm::Impl::copy_back(space, rv, host_staged_rv);
      space.fence();
    });
  }
  req.extend_view_lifetime(sv);
  req.extend_view_lifetime(rv);

  Kokkos::Tools::popRegion();
  return req;
}

I added some unit tests to verify that it somewhat works as is, but I am not yet fully satisfied with the design.

Some open questions:

• Should both functions take an execution space?
• Do we need other functions besides these two?
• Should this PR include GPU-awareness detection of the MPI impl (as given in the example code above)? This may include some non-trivial build system logic to make it work at compile-time since not all MPI impls provide mechanisms to do it then (looking at you MPICH).

[cedricchevalier19]

I think it is not easy to review this PR without the use case.

Do you mind doing another PR with the usage of the API?

[dssgabriel]

Do you mind doing another PR with the usage of the API?

@cedricchevalier19 I have opened #208, which implements host staging for P2P interfaces in the MPI backend. I will do collectives next, but you can start to have a look already.

[dssgabriel]

I am turning this into a draft.

My plan is to rework this PR to let users stage a View in any memory space.
I will close #208 and reimplement host staging solely using create_mirror_view_and_copy.

[cedricchevalier19]

That makes sense, thank you.