use scalar type directly in send/recv

Signed-off-by: Carl Pearson <[email protected]>
kokkos · Jan 29, 2025 · 3b33c77 · 3b33c77
1 parent 1d5d090
commit 3b33c77
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/KokkosComm/mpi/irecv.hpp b/src/KokkosComm/mpi/irecv.hpp
@@ -36,8 +36,8 @@ struct Recv<RecvView, ExecSpace, Mpi> {
     Req<Mpi> req;
     if (KokkosComm::is_contiguous(rv)) {
       space.fence("fence before irecv");
-      MPI_Irecv(KokkosComm::data_handle(rv), 1, view_mpi_type(rv), src, POINTTOPOINT_TAG, h.mpi_comm(),
-                &req.mpi_request());  // TODO: probably best to just use the scalar type
+      MPI_Irecv(KokkosComm::data_handle(rv), KokkosComm::span(rv), mpi_type_v<typename RecvView::value_type>, src,
+                POINTTOPOINT_TAG, h.mpi_comm(), &req.mpi_request());
       req.extend_view_lifetime(rv);
     } else {
       Args args = Packer::allocate_packed_for(space, "TODO", rv);

diff --git a/src/KokkosComm/mpi/isend.hpp b/src/KokkosComm/mpi/isend.hpp
@@ -45,7 +45,8 @@ Req<Mpi> isend_impl(Handle<ExecSpace, Mpi> &h, const SendView &sv, int dest, int
   Req<Mpi> req;
   if (KokkosComm::is_contiguous(sv)) {
     h.space().fence("fence before isend");
-    mpi_isend_fn(KokkosComm::data_handle(sv), 1, view_mpi_type(sv), dest, tag, h.mpi_comm(), &req.mpi_request());
+    mpi_isend_fn(KokkosComm::data_handle(sv), KokkosComm::span(sv), mpi_type_v<typename SendView::value_type>, dest,
+                 tag, h.mpi_comm(), &req.mpi_request());
     req.extend_view_lifetime(sv);
   } else {
     using Packer = typename KokkosComm::PackTraits<SendView>::packer_type;