feat(mpi): implement host staging

[dssgabriel]

This PR is based on #205, using the proposed host staging API.
It enables automatic host staging for the MPI backend, when the provided MPI implementation is not GPU-aware (controlled via a CMake option defined at config-time: -DKokkosComm_ENABLE_GPU_AWARE_MPI=ON).

To-Do list of interfaces to cover

P2P:

• mpi::send
• mpi::isend
• mpi::recv
• mpi::irecv

Colls:

• mpi::broadcast
• mpi::ibroadcast
• mpi::allgather
• mpi::iallgather
• mpi::allreduce
• mpi::iallreduce
• mpi::reduce
• mpi::ireduce (depends on feat(mpi): add non-blocking reduce (mpi:ireduce) #198)
• mpi::alltoall
• mpi::ialltoall
• mpi::inclusive_scan
• mpi::exclusive_scan

Some questions/notes about host staging implementation:

• For non-blocking "send-only" interfaces, do we need to extend the lifetime of the passed view? Technically, only the host-staged view needs to have its lifetime extended until the completion of the communication, not the provided view. E.g.:

  auto host_sv = KokkosComm::Impl::stage_for(sv);
  h.space().fence("fence host staging before `MPI_Isend`");
  
  // Assume contiguous view
  MPI_Isend(data_handle(host_sv), span(host_sv), datatype<MpiSpace, T>(), dest, tag, h.mpi_comm(), &req.mpi_request());
  req.extend_view_lifetime(host_sv);
  
  // NOTE: Is this really needed?
  req.extend_view_lifetime(sv);

• For non-contiguous "receive" interfaces, can we directly unpack into the passed view instead of the host-staged view? This would remove a call to deep_copy, which I think is smart enough to do the right thing, but I am not sure. E.g.:

  auto host_rv = KokkosComm::Impl::stage_for(rv);
  space.fence("fence host staging before `MPI_Recv`");
  
  // Assume non-contiguous view 
  auto packed = Packer::allocate_packed_for(space, "packed `MPI_Recv`", host_rv);
  space.fence("fence packing before `MPI_Recv`");
  MPI_Recv(data_handle(packed.view), packed.count, packed.datatype, src, tag, comm, MPI_STATUS_IGNORE);
  
  // NOTE: Can we unpack directly into `rv` instead of `host_rv`
  // and eliminate the subsequent call to `copy_back`?
  Packer::unpack_into(space, host_rv, args.view);
  KokkosComm::Impl::copy_back(space, rv, host_rv);
  
  space.fence("fence copy back after `MPI_Recv`");

[cedricchevalier19]

Just some thoughts.

To me the biggest issue with staging is how we manage the memory (the same goes for packing).

I am not convinced that we have to wrap Kokkos mirror functions or deep_copy.

[cedricchevalier19]

Not sure if calling create_mirror_view_and_copy and create_mirror_view explicitly is not better

[cedricchevalier19]

Check if copy_back is needed or if we can directly use deep_copy (that should be no-op when the two views are the same).

[dssgabriel]

Yes, this corresponds to my question in the PR description:

• For non-contiguous "receive" interfaces, can we directly unpack into the passed view instead of the host-staged view? This would remove a call to deep_copy, which I think is smart enough to do the right thing, but I am not sure. E.g.:

auto host_rv = KokkosComm::Impl::stage_for(rv);
space.fence("fence host staging before `MPI_Recv`");

// Assume non-contiguous view 
auto packed = Packer::allocate_packed_for(space, "packed `MPI_Recv`", host_rv);
space.fence("fence packing before `MPI_Recv`");
MPI_Recv(data_handle(packed.view), packed.count, packed.datatype, src, tag, comm, MPI_STATUS_IGNORE);

// NOTE: Can we unpack directly into `rv` instead of `host_rv`
// and eliminate the subsequent call to `copy_back`?
Packer::unpack_into(space, host_rv, args.view);
KokkosComm::Impl::copy_back(space, rv, host_rv);

space.fence("fence copy back after `MPI_Recv`");

I'll refactor it with a direct deep_copy to avoid the (unnecessary) intermediate operation.

[cedricchevalier19]

Why we need this?

[dssgabriel]

We don't in the case of a host-staged send operation, since the view that is actually sent (and that needs to live long enough) is host_sv, not sv.
This could be safely removed, but our docs should clearly state what the semantics of KC calls are with respect to view reuse.

In the host-staged case, while sv is technically reusable by the user immediately after the KC collective is called, I think it would be better to have the same semantics in both execution paths, and to mandate that sv is reusable only after the comm operation completes (via wait, wait_all, wait_any, test, etc.).
This also aligns with MPI semantics w.r.t. non-blocking operations.