Merge pull request #30 from dssgabriel/comm-modes

Add communication modes to specialize send routines

docs/api/core.rst

@@ -8,46 +8,60 @@ Core
     * - MPI
       - ``KokkosComm::``
       - ``Kokkos::View``
-    * - MPI_Send
-      - send
+    * - ``MPI_Send``
+      - ``send`` or ``send<CommMode::Standard>``
       - ✓
-    * - MPI_Recv
-      - recv
+    * - ``MPI_Rsend``
+      - ``send<CommMode::Ready>``
       - ✓
-    * - MPI_Isend
-      - isend
+    * - ``MPI_Recv``
+      - ``recv``
       - ✓
-    * - MPI_Reduce
-      - reduce
+    * - ``MPI_Ssend``
+      - ``send<CommMode::Synchronous>``
+      - ✓
+    * - ``MPI_Isend``
+      - ``isend`` or ``isend<CommMode::Standard>``
+      - ✓
+    * - ``MPI_Irsend``
+      - ``isend<CommMode::Ready>``
+      - ✓
+    * - ``MPI_Issend``
+      - ``isend<CommMode::Synchronous>``
+      - ✓
+    * - ``MPI_Reduce``
+      - ``reduce``
       - ✓
 
 Point-to-point
 --------------
 
-.. cpp:function:: template <KokkosExecutionSpace ExecSpace, KokkosView SendView> \
+.. cpp:function:: template <KokkosComm::CommMode SendMode, KokkosExecutionSpace ExecSpace, KokkosView SendView> \
                   Req KokkosComm::isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm)
 
-    MPI_Isend wrapper
+    Wrapper for ``MPI_Isend``, ``MPI_Irsend`` and ``MPI_Issend``.
 
     :param space: The execution space to operate in
     :param sv: The data to send
     :param dest: the destination rank
     :param tag: the MPI tag
     :param comm: the MPI communicator
+    :tparam SendMode: A CommMode_ to use. If unspecified, defaults to a synchronous ``MPI_Issend`` if ``KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE`` is defined, otherwise defaults to a standard ``MPI_Isend``.
     :tparam SendView: A Kokkos::View to send
     :tparam ExecSpace: A Kokkos execution space to operate in
     :returns: A KokkosComm::Req representing the asynchronous communication and any lifetime-extended views.
 
-.. cpp:function:: template <KokkosExecutionSpace ExecSpace, KokkosView SendView> \
+.. cpp:function:: template <KokkosComm::CommMode SendMode, KokkosExecutionSpace ExecSpace, KokkosView SendView> \
                   void KokkosComm::send(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm)
 
-    MPI_Send wrapper
+    Wrapper for ``MPI_Send``, ``MPI_Rsend`` and ``MPI_Ssend``.
 
     :param space: The execution space to operate in
     :param sv: The data to send
     :param dest: the destination rank
     :param tag: the MPI tag
     :param comm: the MPI communicator
+    :tparam SendMode: A CommMode_ to use. If unspecified, defaults to a synchronous ``MPI_Ssend`` if ``KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE`` is defined, otherwise defaults to a standard ``MPI_Send``.
     :tparam SendView: A Kokkos::View to send
     :tparam ExecSpace: A Kokkos execution space to operate in
 
@@ -86,6 +100,29 @@ Collective
 Related Types
 -------------
 
+.. _CommMode:
+
+.. cpp:enum-class:: KokkosComm::CommMode
+
+    A scoped enum to specify the mode of an operation. Buffered mode is not supported.
+
+    .. cpp:enumerator:: KokkosComm::CommMode::Standard
+
+      Standard mode: the MPI implementation decides whether outgoing messages will be buffered. Send operations can be started whether or not a matching receive has been started. They may complete before a matching receive is started. Standard mode is non-local: successful completion of the send operation may depend on the occurrence of a matching receive.
+
+    .. cpp:enumerator:: KokkosComm::CommMode::Ready
+
+      Ready mode: Send operations may be started only if the matching receive is already started.
+
+    .. cpp:enumerator:: KokkosComm::CommMode::Synchronous
+
+      Synchronous mode: Send operations complete successfully only if a matching receive is started, and the receive operation has started to receive the message sent.
+
+    .. cpp:enumerator:: KokkosComm::CommMode::Default
+
+      Default mode is an alias for ``Standard`` mode, but lets users override the behavior of operations at compile-time using the ``KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE`` pre-processor define. This forces ``Synchronous`` mode for all "default-mode" operations, which can be useful for debugging purposes, e.g., for asserting that the communication scheme is correct.
+
+
 .. cpp:class:: KokkosComm::Req
 
     A wrapper around an MPI_Request that can also extend the lifetime of Views.

src/KokkosComm.hpp

@@ -22,21 +22,24 @@
 #include "KokkosComm_recv.hpp"
 #include "KokkosComm_send.hpp"
 #include "KokkosComm_concepts.hpp"
+#include "KokkosComm_comm_mode.hpp"
 
 #include <Kokkos_Core.hpp>
 
 namespace KokkosComm {
 
-template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
+template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace,
+          KokkosView SendView>
 Req isend(const ExecSpace &space, const SendView &sv, int dest, int tag,
           MPI_Comm comm) {
-  return Impl::isend(space, sv, dest, tag, comm);
+  return Impl::isend<SendMode>(space, sv, dest, tag, comm);
 }
 
-template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
+template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace,
+          KokkosView SendView>
 void send(const ExecSpace &space, const SendView &sv, int dest, int tag,
           MPI_Comm comm) {
-  return Impl::send(space, sv, dest, tag, comm);
+  return Impl::send<SendMode>(space, sv, dest, tag, comm);
 }
 
 template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>

src/KokkosComm_comm_mode.hpp

@@ -0,0 +1,43 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#pragma once
+
+namespace KokkosComm {
+
+// Scoped enumeration to specify the communication mode of a sending operation.
+// See section 3.4 of the MPI standard for a complete specification.
+enum class CommMode {
+  // Default mode: lets the user override the send operations behavior at
+  // compile-time. E.g., this can be set to mode "Synchronous" for debug
+  // builds by defining KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE.
+  Default,
+  // Standard mode: MPI implementation decides whether outgoing messages will
+  // be buffered. Send operations can be started whether or not a matching
+  // receive has been started. They may complete before a matching receive is
+  // started. Standard mode is non-local: successful completion of the send
+  // operation may depend on the occurrence of a matching receive.
+  Standard,
+  // Ready mode: Send operations may be started only if the matching receive is
+  // already started.
+  Ready,
+  // Synchronous mode: Send operations complete successfully only if a matching
+  // receive is started, and the receive operation has started to receive the
+  // message sent.
+  Synchronous,
+};
+
+}  // namespace KokkosComm

src/impl/KokkosComm_isend.hpp

@@ -24,13 +24,15 @@
 #include "KokkosComm_pack_traits.hpp"
 #include "KokkosComm_request.hpp"
 #include "KokkosComm_traits.hpp"
+#include "KokkosComm_comm_mode.hpp"
 
 // impl
 #include "KokkosComm_include_mpi.hpp"
 
 namespace KokkosComm::Impl {
 
-template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
+template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace,
+          KokkosView SendView>
 KokkosComm::Req isend(const ExecSpace &space, const SendView &sv, int dest,
                       int tag, MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Impl::isend");
@@ -40,20 +42,42 @@ KokkosComm::Req isend(const ExecSpace &space, const SendView &sv, int dest,
   using KCT  = KokkosComm::Traits<SendView>;
   using KCPT = KokkosComm::PackTraits<SendView>;
 
+  auto mpi_isend_fn = [](void *mpi_view, int mpi_count,
+                         MPI_Datatype mpi_datatype, int mpi_dest, int mpi_tag,
+                         MPI_Comm mpi_comm, MPI_Request *mpi_req) {
+    if constexpr (SendMode == CommMode::Standard) {
+      MPI_Isend(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm,
+                mpi_req);
+    } else if constexpr (SendMode == CommMode::Ready) {
+      MPI_Irsend(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm,
+                 mpi_req);
+    } else if constexpr (SendMode == CommMode::Synchronous) {
+      MPI_Issend(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm,
+                 mpi_req);
+    } else if constexpr (SendMode == CommMode::Default) {
+#ifdef KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE
+      MPI_Issend(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm,
+                 mpi_req);
+#else
+      MPI_Isend(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm,
+                mpi_req);
+#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();
-
-    MPI_Isend(KCT::data_handle(args.view), args.count, args.datatype, dest, tag,
-              comm, &req.mpi_req());
+    mpi_isend_fn(KCT::data_handle(args.view), args.count, args.datatype, dest,
+                 tag, comm, &req.mpi_req());
     req.keep_until_wait(args.view);
   } else {
     using SendScalar = typename SendView::value_type;
-    MPI_Isend(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>, dest,
-              tag, comm, &req.mpi_req());
+    mpi_isend_fn(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>,
+                 dest, tag, comm, &req.mpi_req());
     if (KCT::is_reference_counted()) {
       req.keep_until_wait(sv);
     }

src/impl/KokkosComm_send.hpp

@@ -19,27 +19,49 @@
 #include <Kokkos_Core.hpp>
 
 #include "KokkosComm_pack_traits.hpp"
+#include "KokkosComm_comm_mode.hpp"
 
 // impl
 #include "KokkosComm_include_mpi.hpp"
 
 namespace KokkosComm::Impl {
-template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
+
+template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace,
+          KokkosView SendView>
 void send(const ExecSpace &space, const SendView &sv, int dest, int tag,
           MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Impl::send");
 
   using Packer = typename KokkosComm::PackTraits<SendView>::packer_type;
 
+  auto mpi_send_fn = [](void *mpi_view, int mpi_count,
+                        MPI_Datatype mpi_datatype, int mpi_dest, int mpi_tag,
+                        MPI_Comm mpi_comm) {
+    if constexpr (SendMode == CommMode::Standard) {
+      MPI_Send(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm);
+    } else if constexpr (SendMode == CommMode::Ready) {
+      MPI_Rsend(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm);
+    } else if constexpr (SendMode == CommMode::Synchronous) {
+      MPI_Ssend(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm);
+    } else if constexpr (SendMode == CommMode::Default) {
+#ifdef KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE
+      MPI_Ssend(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm);
+#else
+      MPI_Send(mpi_view, mpi_count, mpi_datatype, mpi_dest, mpi_tag, mpi_comm);
+#endif
+    }
+  };
+
   if (KokkosComm::PackTraits<SendView>::needs_pack(sv)) {
     auto args = Packer::pack(space, sv);
     space.fence();
-    MPI_Send(args.view.data(), args.count, args.datatype, dest, tag, comm);
+    mpi_send_fn(args.view.data(), args.count, args.datatype, dest, tag, comm);
   } else {
     using SendScalar = typename SendView::value_type;
-    MPI_Send(sv.data(), sv.span(), mpi_type_v<SendScalar>, dest, tag, comm);
+    mpi_send_fn(sv.data(), sv.span(), mpi_type_v<SendScalar>, dest, tag, comm);
   }
 
   Kokkos::Tools::popRegion();
 }
+
 }  // namespace KokkosComm::Impl