Merge pull request #109 from cwpearson/feature/transport

Reorganize for MPI / NCCL transport layers

Author: cwpearson

.github/workflows/docs-deploy.yml

@@ -3,7 +3,6 @@
 # To get started with mdBook see: https://rust-lang.github.io/mdBook/index.html
 #
 name: Deploy sphinx site to Pages
-run-name: Deploy sphinx site to Pages
 
 on:
   # Runs on pushes targeting the default branch

.github/workflows/linux-compileonly.yaml

@@ -66,5 +66,5 @@ jobs:
           cmake --build "$KOKKOS_BUILD" --parallel $(nproc) -t install
       - name: Build Kokkos Comm
         run: |
-          cmake -S "$COMM_SRC" -B "$COMM_BUILD" -DCMAKE_CXX_COMPILER="$KOKKOS_SRC/bin/nvcc_wrapper" -DKokkos_ROOT="$KOKKOS_INSTALL" -DCMAKE_BUILD_TYPE=Release
+          cmake -S "$COMM_SRC" -B "$COMM_BUILD" -DCMAKE_CXX_COMPILER="$KOKKOS_SRC/bin/nvcc_wrapper" -DKokkos_ROOT="$KOKKOS_INSTALL" -DCMAKE_BUILD_TYPE=Release -DKokkosComm_ENABLE_TESTS=ON -DKokkosComm_ENABLE_PERFTESTS=ON
           VERBOSE=1 cmake --build "$COMM_BUILD"

.github/workflows/linux.yaml

@@ -62,7 +62,7 @@ jobs:
       COMM_SRC: ${{ github.workspace }}
       COMM_BUILD: ${{ github.workspace }}/build
     runs-on: ubuntu-latest
-    timeout-minutes: 5
+    timeout-minutes: 10
     steps:
       - name: Install MPI
         run: |

CMakeLists.txt

@@ -6,11 +6,13 @@ project(KokkosComm VERSION 0.0.2)
 
 option(KokkosComm_ENABLE_PERFTESTS "Build KokkosComm perf tests" OFF)
 option(KokkosComm_ENABLE_TESTS "Build KokkosComm perf tests" OFF)
+option(KokkosComm_ENABLE_MPI "Build KokkosComm with MPI transport" ON)
 
 
 ## resolve options
 set(KOKKOSCOMM_ENABLE_PERFTESTS ${KokkosComm_ENABLE_PERFTESTS} CACHE BOOL "" FORCE)
 set(KOKKOSCOMM_ENABLE_TESTS ${KokkosComm_ENABLE_TESTS} CACHE BOOL "" FORCE)
+set(KOKKOSCOMM_ENABLE_MPI ${KokkosComm_ENABLE_MPI} CACHE BOOL "" FORCE)
 
 find_package(Kokkos REQUIRED)
 find_package(MPI REQUIRED)

cmake/Config.cmake.in

@@ -8,6 +8,8 @@ include(CMakeFindDependencyMacro)
 find_dependency(MPI)
 find_dependency(Kokkos)
 
+set(KOKKOSCOMM_ENABLE_MPI @KOKKOSCOMM_ENABLE_MPI@)
+
 ## FIXME: do we need this?
 set(KokkosComm_INCLUDE_DIR   "@CMAKE_INSTALL_FULL_INCLUDEDIR@" )
 set(KokkosComm_DATA_DIR   "@CMAKE_INSTALL_PREFIX@/@RELATIVE_DATA_INSTALL_DIR@" )

cmake/KokkosComm_config.hpp.in

@@ -19,3 +19,5 @@
 #define KOKKOSCOMM_VERSION_MAJOR @KOKKOSCOMM_VERSION_MAJOR@
 #define KOKKOSCOMM_VERSION_MINOR @KOKKOSCOMM_VERSION_MINOR@
 #define KOKKOSCOMM_VERSION_PATCH @KOKKOSCOMM_VERSION_PATCH@
+
+#cmakedefine KOKKOSCOMM_ENABLE_MPI

docs/CONTRIBUTING.rst

@@ -27,7 +27,7 @@ Alternatively, you can use docker/podman: (expects $PWD to be the kokkos-comm tr
 .. code-block:: bash
 
   shopt -s globstar
-  podman run -v $PWD:/src xianpengshen/clang-tools:14 clang-format -i {src,unit_tests,perf_tests}/**/*.[ch]pp
+  podman run --rm -v ${PWD}:/src ghcr.io/cwpearson/clang-format-14 clang-format -i {src,unit_tests,perf_tests}/**/*.[ch]pp
 
 Site-Specific Documentation
 ---------------------------

docs/api/core.rst

@@ -1,162 +1,118 @@
 Core
 ====
 
-.. list-table:: MPI API Support
-    :widths: 40 30 15
-    :header-rows: 1
-
-    * - MPI
-      - ``KokkosComm::``
-      - ``Kokkos::View``
-    * - ``MPI_Send``
-      - ``send`` or ``send(KokkosComm::DefaultCommMode{}, ...)``
-      - ✓
-    * - ``MPI_Rsend``
-      - ``send(KokkosComm::ReadyCommMode{}, ...)``
-      - ✓
-    * - ``MPI_Recv``
-      - ``recv``
-      - ✓
-    * - ``MPI_Ssend``
-      - ``send(KokkosComm::SynchronousCommMode{}, ...)``
-      - ✓
-    * - ``MPI_Isend``
-      - ``isend`` or ``isend(KokkosComm::DefaultCommMode{}, ...)``
-      - ✓
-    * - ``MPI_Irsend``
-      - ``isend(KokkosComm::ReadyCommMode{}, ...)``
-      - ✓
-    * - ``MPI_Issend``
-      - ``isend(KokkosComm::SynchronousCommMode{}, ...)``
-      - ✓
-    * - ``MPI_Reduce``
-      - ``reduce``
-      - ✓
-
 Point-to-point
 --------------
 
-.. cpp:function:: template <CommunicationMode SendMode, KokkosExecutionSpace ExecSpace, KokkosView SendView> \
-                  Req KokkosComm::isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm)
+.. cpp:namespace:: KokkosComm
 
-    Wrapper for ``MPI_Isend``, ``MPI_Irsend`` and ``MPI_Issend``.
+.. cpp:function:: template <KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, Transport TRANSPORT = DefaultTransport> Req<TRANSPORT> send(Handle<ExecSpace, TRANSPORT> &h, SendView &sv, int dest)
 
-    :param mode: The communication mode to use
-    :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 IsendMode: A communication mode to use, one of: ``KokkosComm::DefaultCommMode``, ``KokkosComm::StandardCommMode``, ``KokkosComm::SynchronousCommMode`` or ``KokkosComm::ReadyCommMode`` (modeled with the ``KokkosComm::CommunicationMode`` concept)
-    :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.
+  Initiates a non-blocking send operation.
 
-.. cpp:function:: template <typename SendMode, KokkosExecutionSpace ExecSpace, KokkosView SendView> \
-                  void KokkosComm::send(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm)
+  .. warning::
+    This is not a blocking operation despite being named like ``MPI_Send``.
 
-    Wrapper for ``MPI_Send``, ``MPI_Rsend`` and ``MPI_Ssend``.
+  :tparam SendView: The type of the Kokkos view to send.
+  :tparam ExecSpace: The execution space to use. Defaults to Kokkos::DefaultExecutionSpace.
+  :tparam TRANSPORT: The transport mechanism to use. Defaults to DefaultTransport.
 
-    :param mode: The communication mode to use
-    :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 communication mode to use, one of: ``KokkosComm::DefaultCommMode``, ``KokkosComm::StandardCommMode``, ``KokkosComm::SynchronousCommMode`` or ``KokkosComm::ReadyCommMode`` (modeled with the ``KokkosComm::CommunicationMode`` concept)
-    :tparam SendView: A Kokkos::View to send
-    :tparam ExecSpace: A Kokkos execution space to operate in
+  :param h: A handle to the execution space and transport mechanism.
+  :param sv: The Kokkos view to send.
+  :param dest: The destination rank.
 
-.. cpp:function:: template <KokkosExecutionSpace ExecSpace, KokkosView RecvView> \
-                  void KokkosComm::recv(const ExecSpace &space, RecvView &rv, int src, int tag, MPI_Comm comm)
+  :return: A request object for the non-blocking send operation.
 
-    MPI_Recv wrapper
+.. cpp:function:: template <KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, Transport TRANSPORT = DefaultTransport> Req<TRANSPORT> send(SendView &sv, int dest)
 
-    :param space: The execution space to operate in
-    :param srv: The data to recv
-    :param src: the source rank
-    :param tag: the MPI tag
-    :param comm: the MPI communicator
-    :tparam Recv: A Kokkos::View to send
-    :tparam ExecSpace: A Kokkos execution space to operate in
+   Initiates a non-blocking send operation using a default handle.
 
+   .. warning::
+     This is not a blocking operation despite being named like ``MPI_Send``.
 
-Collective
-----------
+   :tparam SendView: The type of the Kokkos view to send.
+   :tparam ExecSpace: The execution space to use. Defaults to Kokkos::DefaultExecutionSpace.
+   :tparam TRANSPORT: The transport mechanism to use. Defaults to DefaultTransport.
 
-.. cpp:function:: template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvView> \
-                  void KokkosComm::reduce(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Op op, int root, MPI_Comm comm)
+   :param sv: The Kokkos view to send.
+   :param dest: The destination rank.
 
-    MPI_Reduce wrapper
+   :return: A request object for the non-blocking send operation.
 
-    :param space: The execution space to operate in
-    :param sv: The data to send
-    :param rv: The view to receive into
-    :param op: The MPI_Op to use in the reduction
-    :param root: The root rank for the reduction
-    :param comm: the MPI communicator
-    :tparam SendView: A Kokkos::View to send
-    :tparam RecvView: A Kokkos::View to recv
-    :tparam ExecSpace: A Kokkos execution space to operate in
+   Example usage:
 
+.. literalinclude:: core_send.cpp
+   :language: cpp
 
-.. cpp:function:: template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvView> \
-                  void KokkosComm::allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Comm comm)
 
-    MPI_Allgather wrapper
 
-    :param space: The execution space to operate in
-    :param sv: The data to send
-    :param rv: The view to receive into
-    :param comm: the MPI communicator
-    :tparam SendView: A Kokkos::View to send. Contiguous and rank less than 2.
-    :tparam RecvView: A Kokkos::View to recv. Contiguous and rank 1.
-    :tparam ExecSpace: A Kokkos execution space to operate in
+.. cpp:function:: template <KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, Transport TRANSPORT = DefaultTransport> Req<TRANSPORT> recv(Handle<ExecSpace, TRANSPORT> &h, RecvView &rv, int src)
 
-    If ``sv`` is a rank-0 view, the value from the jth rank will be placed in index j of ``rv``.
+   Initiates a non-blocking receive operation.
 
-Related Types
--------------
+   .. warning::
+     This is not a blocking operation despite being named like ``MPI_Recv``.
+
+   :tparam RecvView: The type of the Kokkos view for receiving data.
+   :tparam ExecSpace: The execution space where the operation will be performed. Defaults to `Kokkos::DefaultExecutionSpace`.
+   :tparam TRANSPORT: The transport mechanism to be used. Defaults to `DefaultTransport`.
+
+   :param h: A handle to the execution space and transport mechanism.
+   :param rv: The Kokkos view where the received data will be stored.
+   :param src: The source rank from which to receive data.
+
+   :return: A request object of type `Req<TRANSPORT>` representing the non-blocking receive operation.
+
+   This function initiates a non-blocking receive operation using the specified execution space and transport mechanism. The data will be received into the provided view from the specified source rank and message tag. The function returns a request object that can be used to check the status of the receive operation or to wait for its completion.
 
-Communication Modes
-^^^^^^^^^^^^^^^^^^^
+   Example usage:
 
-Structures to specify the mode of an operation. Buffered mode is not supported.
+.. literalinclude:: core_recv.cpp
+   :language: cpp
 
-.. cpp:struct:: KokkosComm::StandardCommMode
 
-  Let the MPI implementation decide 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 begins. Standard mode is non-local: successful completion of the send operation may depend on the occurrence of a matching receive.
 
-.. cpp:struct:: KokkosComm::SynchronousCommMode
 
-  Send operations complete successfully only if a matching receive is started, and the receive operation has started to receive the message sent.
 
-.. cpp:struct:: KokkosComm::ReadyCommMode
+.. cpp:function:: template <KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, Transport TRANSPORT = DefaultTransport> Req<TRANSPORT> recv(RecvView &rv, int src)
 
-  Send operations may be started only if the matching receive is already started.
+   Initiates a non-blocking receive operation using a default handle.
 
-.. cpp:struct:: KokkosComm::DefaultCommMode
+   .. warning::
+     This is not a blocking operation despite being named like ``MPI_Recv``.
 
-  The default mode is aliased as ``Standard`` but lets users override the behavior of operations at compile-time using the ``KOKKOSCOMM_FORCE_SYNCHRONOUS_MODE`` pre-processor definition. The latter forces ``Synchronous`` mode for all "default-mode" operations, which can be helpful for debugging purposes, e.g., asserting that the communication scheme is correct.
+   :tparam RecvView: The type of the Kokkos view for receiving data.
+   :tparam ExecSpace: The execution space where the operation will be performed. Defaults to `Kokkos::DefaultExecutionSpace`.
+   :tparam TRANSPORT: The transport mechanism to be used. Defaults to `DefaultTransport`.
 
+   :param rv: The Kokkos view where the received data will be stored.
+   :param src: The source rank from which to receive data.
 
-Requests
-^^^^^^^^
+   :return: A request object of type `Req<TRANSPORT>` representing the non-blocking receive operation.
 
-.. cpp:class:: KokkosComm::Req
 
-    A wrapper around an MPI_Request that can also extend the lifetime of Views.
+Collective
+----------
+
+.. cpp:namespace:: KokkosComm
+
+.. cpp:function:: template <KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace, Transport TRANSPORT = DefaultTransport> void barrier(Handle<ExecSpace, TRANSPORT> &&h)
+
+   A function to create a barrier using the given execution space and transport handle.
 
-    .. cpp:function:: MPI_Request &KokkosComm::Req::mpi_req()
+   :tparam ExecSpace: The execution space to be used. Defaults to `Kokkos::DefaultExecutionSpace`.
+   :tparam TRANSPORT: The transport mechanism to be used. Defaults to `DefaultTransport`.
+   :param h: A handle of type `Handle<ExecSpace, TRANSPORT>` to be forwarded to the barrier implementation.
 
-        Retrieve a reference to the held MPI_Request.
 
-    .. cpp:function:: void KokkosComm::Req::wait()
 
-        Call MPI_Wait on the held MPI_Request and drop copies of any previous arguments to Req::keep_until_wait().
+Related Types
+-------------
+
+.. cpp:namespace:: KokkosComm
+
+.. cpp:class:: template <Transport TRANSPORT = DefaultTransport> Req
 
-    .. cpp:function:: template<typename View> \
-                      void KokkosComm::Req::keep_until_wait(const View &v)
+   A template class to handle requests with different transport types.
 
-        Extend the lifetime of v at least until Req::wait() is called.
-        This is useful to prevent a View from being destroyed during an asynchronous MPI operation.
+   :tparam TRANSPORT: The type of transport. Defaults to :cpp:enumerator:`KokkosComm::DefaultTransport`.

docs/api/core_recv.cpp

@@ -0,0 +1,4 @@
+Handle<> handle;
+Kokkos::View<double*> recv_view("recv_view", 100);
+auto req = recv(handle, recv_view, 1/*src*/);
+KokkosComm::wait(req);

docs/api/core_send.cpp

@@ -0,0 +1,29 @@
+#include <Kokkos_Comm.hpp>
+
+// Define the execution space and transport
+using ExecSpace = Kokkos::DefaultExecutionSpace;
+using Transport = DefaultTransport;
+
+// Create a Kokkos view
+Kokkos::View<double*> data("data", 100);
+
+// Fill the view with some data
+Kokkos::parallel_for("fill_data", Kokkos::RangePolicy<ExecSpace>(0, 100), KOKKOS_LAMBDA(int i) {
+  data(i) = static_cast<double>(i);
+});
+
+// Destination rank and message tag
+int dest = 1;
+
+// Create a handle
+KokkosComm::Handle<> handle; // Same as Handle<Execspace, Transport>
+
+// Initiate a non-blocking send with a handle
+auto req1 = send(handle, data, dest);
+
+// Initiate a non-blocking send with a default handle
+auto req2 = send(data, dest);
+
+// Wait for the requests to complete (assuming a wait function exists)
+KokkosComm::wait(req1);
+KokkosComm::wait(req2);