Added alltoall wrapper, revised RecvView variables to rv

kokkos · Jun 12, 2024 · a9e03db · a9e03db
1 parent 5e17145
commit a9e03db
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/src/KokkosComm.hpp b/src/KokkosComm.hpp
@@ -22,6 +22,7 @@
 #include "KokkosComm_irecv.hpp"
 #include "KokkosComm_recv.hpp"
 #include "KokkosComm_send.hpp"
+#include "KokkosComm_alltoall.hpp"
 #include "KokkosComm_concepts.hpp"
 #include "KokkosComm_comm_mode.hpp"
 
@@ -35,8 +36,8 @@ Req isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Com
 }
 
 template <KokkosView RecvView>
-void irecv(RecvView &sv, int src, int tag, MPI_Comm comm, MPI_Request req) {
-  return Impl::irecv(sv, src, tag, comm, req);
+void irecv(RecvView &rv, int src, int tag, MPI_Comm comm, MPI_Request req) {
+  return Impl::irecv(rv, src, tag, comm, req);
 }
 
 template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace, KokkosView SendView>
@@ -45,8 +46,14 @@ void send(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Com
 }
 
 template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
-void recv(const ExecSpace &space, RecvView &sv, int src, int tag, MPI_Comm comm) {
-  return Impl::recv(space, sv, src, tag, comm);
+void recv(const ExecSpace &space, RecvView &rv, int src, int tag, MPI_Comm comm) {
+  return Impl::recv(space, rv, src, tag, comm);
+}
+
+template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvView>
+void alltoall(const ExecSpace &space, const SendView &sv, const size_t sendCount, const RecvView &rv,
+              const size_t recvCount, MPI_Comm comm) {
+  return Impl::alltoall(space, sv, sendCount, rv, recvCount, comm);
 }
 
 }  // namespace KokkosComm