refactor: define a lambda to avoid logic duplication in pack vs. non-pack paths

Author: dssgabriel

src/impl/KokkosComm_isend.hpp

@@ -42,52 +42,35 @@ KokkosComm::Req isend(const ExecSpace &space, const SendView &sv, int dest,
   using KCT  = KokkosComm::Traits<SendView>;
   using KCPT = KokkosComm::PackTraits<SendView>;
 
-  if (KCPT::needs_pack(sv)) {
-    using Packer  = typename KCPT::packer_type;
-    using MpiArgs = typename Packer::args_type;
-
-    MpiArgs args = Packer::pack(space, sv);
-    space.fence();
-
+  auto isend_fn = [&](auto &&view, auto count, auto datatype, auto &&request) {
     if constexpr (SendMode == CommMode::Standard) {
-      MPI_Isend(KCT::data_handle(args.view), args.count, args.datatype, dest,
-                tag, comm, &req.mpi_req());
+      MPI_Isend(view, count, datatype, dest, tag, comm, request);
     } else if constexpr (SendMode == CommMode::Ready) {
-      MPI_Irsend(KCT::data_handle(args.view), args.count, args.datatype, dest,
-                 tag, comm, &req.mpi_req());
+      MPI_Irsend(view, count, datatype, dest, tag, comm, request);
     } else if constexpr (SendMode == CommMode::Synchronous) {
-      MPI_Issend(KCT::data_handle(args.view), args.count, args.datatype, dest,
-                 tag, comm, &req.mpi_req());
+      MPI_Issend(view, count, datatype, dest, tag, comm, request);
     } else if constexpr (SendMode == CommMode::Default) {
 #ifdef KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE
-      MPI_Issend(KCT::data_handle(args.view), args.count, args.datatype, dest,
-                 tag, comm, &req.mpi_req());
+      MPI_Issend(view, count, datatype, dest, tag, comm, request);
 #else
-      MPI_Isend(KCT::data_handle(args.view), args.count, args.datatype, dest,
-                tag, comm, &req.mpi_req());
+      MPI_Isend(view, count, datatype, dest, tag, comm, request);
 #endif
     }
+  };
+
+  if (KCPT::needs_pack(sv)) {
+    using Packer  = typename KCPT::packer_type;
+    using MpiArgs = typename Packer::args_type;
+
+    MpiArgs args = Packer::pack(space, sv);
+    space.fence();
+    isend_fn(KCT::data_handle(args.view), args.count, args.datatype,
+             &req.mpi_req());
     req.keep_until_wait(args.view);
   } else {
     using SendScalar = typename SendView::value_type;
-    if constexpr (SendMode == CommMode::Standard) {
-      MPI_Isend(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>,
-                dest, tag, comm, &req.mpi_req());
-    } else if constexpr (SendMode == CommMode::Ready) {
-      MPI_Irsend(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>,
-                 dest, tag, comm, &req.mpi_req());
-    } else if constexpr (SendMode == CommMode::Synchronous) {
-      MPI_Issend(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>,
-                 dest, tag, comm, &req.mpi_req());
-    } else if constexpr (SendMode == CommMode::Default) {
-#ifdef KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE
-      MPI_Issend(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>,
-                 dest, tag, comm, &req.mpi_req());
-#else
-      MPI_Isend(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>,
-                dest, tag, comm, &req.mpi_req());
-#endif
-    }
+    isend_fn(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>,
+             &req.mpi_req());
     if (KCT::is_reference_counted()) {
       req.keep_until_wait(sv);
     }

src/impl/KokkosComm_send.hpp

@@ -34,37 +34,29 @@ void send(const ExecSpace &space, const SendView &sv, int dest, int tag,
 
   using Packer = typename KokkosComm::PackTraits<SendView>::packer_type;
 
-  if (KokkosComm::PackTraits<SendView>::needs_pack(sv)) {
-    auto args = Packer::pack(space, sv);
-    space.fence();
+  auto send_fn = [&](auto &&view, auto count, auto datatype) {
     if constexpr (SendMode == CommMode::Standard) {
-      MPI_Send(args.view.data(), args.count, args.datatype, dest, tag, comm);
+      MPI_Send(view, count, datatype, dest, tag, comm);
     } else if constexpr (SendMode == CommMode::Ready) {
-      MPI_Rsend(args.view.data(), args.count, args.datatype, dest, tag, comm);
+      MPI_Rsend(view, count, datatype, dest, tag, comm);
     } else if constexpr (SendMode == CommMode::Synchronous) {
-      MPI_Ssend(args.view.data(), args.count, args.datatype, dest, tag, comm);
+      MPI_Ssend(view, count, datatype, dest, tag, comm);
     } else if constexpr (SendMode == CommMode::Default) {
 #ifdef KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE
-      MPI_Ssend(args.view.data(), args.count, args.datatype, dest, tag, comm);
+      MPI_Ssend(view, count, datatype, dest, tag, comm);
 #else
-      MPI_Send(args.view.data(), args.count, args.datatype, dest, tag, comm);
+      MPI_Send(view, count, datatype, dest, tag, comm);
 #endif
     }
+  };
+
+  if (KokkosComm::PackTraits<SendView>::needs_pack(sv)) {
+    auto args = Packer::pack(space, sv);
+    space.fence();
+    send_fn(args.view.data(), args.count, args.datatype);
   } else {
     using SendScalar = typename SendView::value_type;
-    if constexpr (SendMode == CommMode::Standard) {
-      MPI_Send(sv.data(), sv.span(), mpi_type_v<SendScalar>, dest, tag, comm);
-    } else if constexpr (SendMode == CommMode::Ready) {
-      MPI_Rsend(sv.data(), sv.span(), mpi_type_v<SendScalar>, dest, tag, comm);
-    } else if constexpr (SendMode == CommMode::Synchronous) {
-      MPI_Ssend(sv.data(), sv.span(), mpi_type_v<SendScalar>, dest, tag, comm);
-    } else if constexpr (SendMode == CommMode::Default) {
-#ifdef KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE
-      MPI_Ssend(sv.data(), sv.span(), mpi_type_v<SendScalar>, dest, tag, comm);
-#else
-      MPI_Send(sv.data(), sv.span(), mpi_type_v<SendScalar>, dest, tag, comm);
-#endif
-    }
+    send_fn(sv.data(), sv.span(), mpi_type_v<SendScalar>);
   }
 
   Kokkos::Tools::popRegion();