kokkos / kokkos-comm

Experimental MPI Wrapper for Kokkos
https://kokkos.org/kokkos-comm/
Other
12 stars 9 forks source link

`KokkosComm_irecv.hpp` has no function overload to take an `ExecSpace &space` parameter #102

Open nicoleavans opened 3 weeks ago

nicoleavans commented 3 weeks ago

Current state of usage:

  if (rank == 0) {
    KokkosComm::Req sendreq = KokkosComm::isend(space, v, 1, 1, comm);
    sendreq.wait();
  } else if (rank == 1) {
    KokkosComm::Req recvreq = KokkosComm::irecv(v, 0, 1, comm);
    recvreq.wait();
  }

Current functions in irecv:

template <KokkosView RecvView>
void irecv(RecvView &rv, int src, int tag, MPI_Comm comm, MPI_Request &req) {
  Kokkos::Tools::pushRegion("KokkosComm::Impl::irecv");

  if (KokkosComm::is_contiguous(rv)) {
    using RecvScalar = typename RecvView::value_type;
    MPI_Irecv(KokkosComm::data_handle(rv), KokkosComm::span(rv), mpi_type_v<RecvScalar>, src, tag, comm, &req);
  } else {
    throw std::runtime_error("Only contiguous irecv viewsupported");
  }

  Kokkos::Tools::popRegion();
}

template <KokkosView RecvView>
KokkosComm::Req irecv(RecvView &rv, int src, int tag, MPI_Comm comm) {
  Kokkos::Tools::pushRegion("KokkosComm::Impl::irecv");
  KokkosComm::Req req;
  irecv(rv, src, tag, comm, req.mpi_req());
  return req;
}