Merge branch 'develop' into pr/rank-supported

Author: cedricchevalier19

.github/workflows/snl-h100.yaml

@@ -1,65 +1,133 @@
 name: H100
 
-
 permissions:
   contents: none
 
 on:
   workflow_call:
 
 jobs:
-  PR_CUDA1250_OPENMPI504:
-    name: PR_CUDA1250_OPENMPI504
-    runs-on: [cuda125-openmpi504-latest-latest]
+  PR_CUDA1262_OPENMPI505:
+    name: PR_CUDA1262_OPENMPI505
 
+    # this label is correct, but the underlying AT2 runner is actually
+    # CUDA 12.6.2 and OpenMPI 5.0.5 <facepalm>
+    runs-on: [cuda125-openmpi504-latest-latest]
     steps:
-      - name: Checkout Kokkos Comm
+      - name: Check NVIDIA GPU
+        run: nvidia-smi
+
+      - name: Kokkos - Checkout
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          repository: kokkos/kokkos
+          ref: 4.7.01
+          path: kokkos
+      - name: Kokkos - Configure
+        run: >
+          cmake
+          -S kokkos
+          -B kokkos/build
+          -DCMAKE_CXX_COMPILER=$(realpath kokkos/bin/nvcc_wrapper)
+          -DCMAKE_CXX_STANDARD=20
+          -DCMAKE_CXX_EXTENSIONS=OFF
+          -DCMAKE_INSTALL_PREFIX=kokkos/install
+          -DKokkos_ENABLE_CUDA=ON
+          -DKokkos_ARCH_HOPPER90=ON
+          -DKokkos_ENABLE_TESTS=OFF
+          -DKokkos_ENABLE_EXAMPLES=OFF
+          -DKokkos_ENABLE_DEPRECATION_WARNINGS=OFF
+          -DKokkos_ENABLE_DEPRECATED_CODE_4=OFF
+          -DKokkos_ENABLE_DEPRECATION_WARNINGS=OFF
+      - name: Kokkos - Build
+        run: cmake --build kokkos/build --parallel $(nproc)
+      - name: Kokkos - Install
+        run: cmake --build kokkos/build --target install --parallel $(nproc)
+      - name: KokkosComm - Checkout
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           path: kokkos-comm
+      - name: KokkosComm - Configure MPI backend
+        run: >
+          cmake
+          -S kokkos-comm
+          -B build-mpi
+          -DCMAKE_CXX_COMPILER=$(realpath kokkos/bin/nvcc_wrapper)
+          -DCMAKE_CXX_STANDARD=20
+          -DCMAKE_CXX_EXTENSIONS=OFF
+          -DCMAKE_CXX_FLAGS="-Werror"
+          -DKokkos_ROOT=kokkos/install
+          -DKokkosComm_ENABLE_MPI=ON
+          -DKokkosComm_ENABLE_TESTS=ON
+          -DKokkosComm_ENABLE_PERFTESTS=ON
+      - name: KokkosComm - Build MPI backend
+        run: cmake --build build-mpi --parallel $(nproc)
+      - name: KokkosComm - Test MPI backend
+        working-directory: build-mpi
+        run: ctest --output-on-failure -V --timeout 1200
 
-      - name: Checkout Kokkos
+  PR_CUDA1262_NCCL2275:
+    name: PR_CUDA1262_NCCL2275
+
+    # this label is correct, but the underlying AT2 runner is actually
+    # CUDA 12.6.2 and OpenMPI 5.0.5 <facepalm>
+    runs-on: [cuda125-openmpi504-latest-latest]
+    steps:
+      - name: Check NVIDIA GPU
+        run: nvidia-smi
+
+      - name: Kokkos - Checkout
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           repository: kokkos/kokkos
-          ref: 4.5.01
+          ref: 4.7.01
           path: kokkos
 
-      - name: nvidia-smi
-        run: nvidia-smi
+      - name: Kokkos - Configure
+        run: >
+          cmake
+          -S kokkos
+          -B kokkos/build
+          -DCMAKE_CXX_COMPILER=$(realpath kokkos/bin/nvcc_wrapper)
+          -DCMAKE_CXX_STANDARD=20
+          -DCMAKE_CXX_EXTENSIONS=OFF
+          -DCMAKE_INSTALL_PREFIX=kokkos/install
+          -DKokkos_ENABLE_CUDA=ON
+          -DKokkos_ARCH_HOPPER90=ON
+          -DKokkos_ENABLE_TESTS=OFF
+          -DKokkos_ENABLE_EXAMPLES=OFF
+          -DKokkos_ENABLE_DEPRECATION_WARNINGS=OFF
+          -DKokkos_ENABLE_DEPRECATED_CODE_4=OFF
+          -DKokkos_ENABLE_DEPRECATION_WARNINGS=OFF
+      - name: Kokkos - Build
+        run: cmake --build kokkos/build --parallel $(nproc)
 
-      - name: Configure Kokkos
-        run: |
-          cmake -S kokkos -B kokkos/build \
-            -DCMAKE_CXX_COMPILER=$(realpath kokkos/bin/nvcc_wrapper) \
-            -DCMAKE_CXX_STANDARD=20 \
-            -DCMAKE_CXX_EXTENSIONS=OFF \
-            -DCMAKE_INSTALL_PREFIX=kokkos/install \
-            -DKokkos_ENABLE_CUDA=ON \
-            -DKokkos_ARCH_HOPPER90=ON \
-            -DKokkos_ENABLE_TESTS=OFF \
-            -DKokkos_ENABLE_EXAMPLES=OFF \
-            -DKokkos_ENABLE_DEPRECATION_WARNINGS=OFF \
-            -DKokkos_ENABLE_DEPRECATED_CODE_4=OFF \
-            -DKokkos_ENABLE_DEPRECATION_WARNINGS=OFF
-
-      - name: Build Kokkos
+      - name: Kokkos - Install
         run: cmake --build kokkos/build --target install --parallel $(nproc)
 
-      - name: Configure Kokkos Comm
-        run: |
-          cmake -S kokkos-comm -B build \
-            -DCMAKE_CXX_COMPILER=$(realpath kokkos/bin/nvcc_wrapper) \
-            -DCMAKE_CXX_STANDARD=20 \
-            -DCMAKE_CXX_EXTENSIONS=OFF \
-            -DCMAKE_CXX_FLAGS="-Werror" \
-            -DKokkos_ROOT=kokkos/install \
-            -DKokkosComm_ENABLE_TESTS=ON \
-            -DKokkosComm_ENABLE_PERFTESTS=ON
+      - name: KokkosComm - Checkout
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          path: kokkos-comm
+
+      - name: KokkosComm - Configure NCCL backend
+        # FIXME_NCCL: no performance tests on NCCL
+        run: >
+          cmake
+          -S kokkos-comm
+          -B build-nccl
+          -DCMAKE_CXX_COMPILER=$(realpath kokkos/bin/nvcc_wrapper)
+          -DCMAKE_CXX_STANDARD=20
+          -DCMAKE_CXX_EXTENSIONS=OFF
+          -DKokkos_ROOT=kokkos/install
+          -DKokkosComm_ENABLE_MPI=OFF
+          -DKokkosComm_ENABLE_NCCL=ON
+          -DKokkosComm_ENABLE_TESTS=ON
+          -DKokkosComm_ENABLE_PERFTESTS=OFF
 
-      - name: Build Kokkos Comm
-        run: cmake --build build --parallel $(nproc)
+      - name: KokkosComm - Build NCCL backend
+        run: cmake --build build-nccl --parallel $(nproc)
 
-      - name: Test Kokkos Comm
-        working-directory: build
+      - name: KokkosComm - Test NCCL backend
+        working-directory: build-nccl
         run: ctest --output-on-failure -V --timeout 1200

docs/design/nccl_interop.rst

@@ -0,0 +1,14 @@
+*********************
+NCCL interoperability
+*********************
+
+There are several challenges with supporting NCCL.
+
+* For multi-process NCCL, we need a way to share a unique ID between processes so that the different processes know they're part of the same NCCL communicator. For the time being, this is accomplished via MPI in the nccl tests.
+
+* NCCL has the concept of a non-blocking communicator. This causes all NCCL operations to potentially return ``ncclInProgress`` BEFORE they actually put an GPU operations into streams. This means we can't just synchronize on a CUDA stream in the NCCL backend's ``wait`` implementation. Either:
+
+  * our NCCL operations need to effectively become blocking (checking NCCL's async status thing until it's no longer in progress)
+  * our ``wait`` implementation needs to do that
+
+* It is not guaranteed to be safe for NCCL operations and GPU-aware MPI operations to be simultaneously active on the same set of GPUs. This is a challenge both if we permit multiple backends to exist, and for interop with existing MPI and/or NCCL applications.

docs/dev/impl_comm_space.rst

@@ -20,9 +20,8 @@ For example, for the MPI communication space, we define the following:
 
     namespace KokkosComm {
 
-    struct Mpi {
-      static auto world_size() noexcept -> int { /* ... */ }
-      static auto world_rank() noexcept -> int { /* ... */ }
+    struct MpiSpace {
+      ...
     };
 
     template <>
@@ -31,7 +30,6 @@ For example, for the MPI communication space, we define the following:
     } // end KokkosComm
 
 
-Notice that ``Mpi`` has two static methods, but that these methods are not required. The main point is that ``struct Mpi`` exists.
 To let core API functions know that your communication space is something KokkosComm can use to dispatch messages, you also need to declare the ``Impl::is_communication_space`` specialization using the ``CommunicationSpace`` concept.
 
 
@@ -71,7 +69,7 @@ For example, for the MPI communication space request, we define the following:
     namespace KokkosComm {
 
     template <>
-    class Req<Mpi> { /* ... */ };
+    class Req<MpiSpace> { /* ... */ };
 
     } // end KokkosComm
 
@@ -103,7 +101,7 @@ For example, for the MPI communication space, we create a partial specialization
     namespace KokkosComm::Impl {
 
     template <KokkosView RecvView, KokkosExecutionSpace ExecSpace>
-    struct Recv<RecvView, ExecSpace, Mpi> { /* ... */ };
+    struct Recv<RecvView, ExecSpace, MpiSpace> { /* ... */ };
 
     } // end KokkosComm::Impl
 

docs/index.rst

@@ -29,6 +29,7 @@ Documentation Content
 
    design/overview
    design/mpi_interop
+   design/nccl_interop
 
 .. toctree::
    :maxdepth: 1

perf_tests/CMakeLists.txt

@@ -18,20 +18,14 @@ endif()
 
 include(FetchContent)
 
-# Avoid warning about DOWNLOAD_EXTRACT_TIMESTAMP in CMake 3.24:
-if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.24.0")
-  cmake_policy(SET CMP0135 NEW)
-endif()
-
+FetchContent_Declare(benchmark
+  GIT_REPOSITORY https://github.com/google/benchmark.git
+  GIT_TAG
+    eddb0241389718a23a42db6af5f0164b6e0139af # v1.9.4
+)
 set(BENCHMARK_ENABLE_TESTING OFF CACHE BOOL "" FORCE)
-FetchContent_Declare(benchmark URL https://github.com/google/benchmark/archive/refs/tags/v1.8.3.zip)
-# FetchContent_MakeAvailable(benchmark) was making install benchmark as well
-# EXCLUDE_FROM_ALL here seems to be the magic
-if(NOT benchmark_POPULATED)
-  FetchContent_Populate(benchmark)
-  add_subdirectory(${benchmark_SOURCE_DIR} ${benchmark_BINARY_DIR} EXCLUDE_FROM_ALL)
-endif()
-unset(BENCHMARK_ENABLE_TESTING)
+set(BENCHMARK_ENABLE_INSTALL OFF CACHE BOOL "" FORCE)
+FetchContent_MakeAvailable(benchmark)
 
 if(KOKKOSCOMM_ENABLE_MPI)
   add_subdirectory(mpi)

src/KokkosComm/collective.hpp

@@ -22,20 +22,11 @@
 
 namespace KokkosComm::Experimental {
 
-/// Copy the `sv` view on the `root` rank to all ranks' `rv` view.
-///
-/// The `sv` view is only used on the `root` rank and ignored for all other ranks.
-template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
-          CommunicationSpace CommSpace = DefaultCommunicationSpace>
-auto broadcast(Handle<ExecSpace, CommSpace>& h, const SendView sv, RecvView rv, int root) -> Req<CommSpace> {
-  return Impl::Broadcast<SendView, RecvView, ExecSpace, CommSpace>::execute(h, sv, rv, root);
-}
-
-/// In-place variant of `broadcast`. Copy the `v` view from the `root` rank to all ranks' `v` view.
+/// Copy the `v` view from the `root` rank to all ranks' `v` view.
 template <KokkosView View, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           CommunicationSpace CommSpace = DefaultCommunicationSpace>
 auto broadcast(Handle<ExecSpace, CommSpace>& h, View v, int root) -> Req<CommSpace> {
-  return Impl::Broadcast<View, View, ExecSpace, CommSpace>::execute(h, v, v, root);
+  return Impl::Broadcast<View, ExecSpace, CommSpace>::execute(h, v, root);
 }
 
 /// Copy the `sv` view from each rank to the `rv` view, receiving data from rank `i` at offset

src/KokkosComm/concepts.hpp

@@ -27,7 +27,15 @@ template <typename T>
 concept KokkosExecutionSpace = Kokkos::is_execution_space_v<T>;
 
 template <typename T>
-concept CommunicationSpace = KokkosComm::Impl::is_communication_space<T>::value;
+concept CommunicationSpace = requires {
+  KokkosComm::Impl::is_communication_space<T>::value;
+  typename T::communication_space;
+  typename T::handle_type;
+  typename T::request_type;
+  typename T::datatype_type;
+  typename T::reduction_op_type;
+  typename T::rank_type;
+};
 
 template <typename T>
 concept ReductionOperator = KokkosComm::Impl::is_reduction_operator<T>::value;

src/KokkosComm/fwd.hpp

@@ -9,18 +9,19 @@
 
 namespace KokkosComm {
 
-// NCCL backend also implicitly declares MPI
 #if defined(KOKKOSCOMM_ENABLE_NCCL)
 namespace Experimental {
-class Nccl;
-}  // namespace Experimental
-class Mpi;
-using DefaultCommunicationSpace  = Experimental::Nccl;
-using FallbackCommunicationSpace = Mpi;
+struct NcclSpace;
+}
+// NCCL backend also declares the MPI space as fallback
+struct MpiSpace;
+
+using DefaultCommunicationSpace  = Experimental::NcclSpace;
+using FallbackCommunicationSpace = MpiSpace;
 #elif defined(KOKKOSCOMM_ENABLE_MPI)
-class Mpi;
-using DefaultCommunicationSpace  = Mpi;
-using FallbackCommunicationSpace = Mpi;
+struct MpiSpace;
+using DefaultCommunicationSpace  = MpiSpace;
+using FallbackCommunicationSpace = MpiSpace;
 #else
 #error at least one communication space must be enabled
 #endif
@@ -47,7 +48,7 @@ struct Send;
 // Collectives are currently experimental functions
 namespace Experimental::Impl {
 
-template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+template <KokkosView View, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           CommunicationSpace CommSpace = DefaultCommunicationSpace>
 struct Broadcast;
 

src/KokkosComm/mpi/allgather.hpp

@@ -8,12 +8,38 @@
 
 #include <KokkosComm/concepts.hpp>
 #include <KokkosComm/traits.hpp>
+#include "mpi_space.hpp"
+#include "req.hpp"
 
 #include "impl/types.hpp"
 #include "impl/error_handling.hpp"
 
 namespace KokkosComm::mpi {
 
+template <KokkosExecutionSpace ExecSpace, KokkosView SView, KokkosView RView>
+auto iallgather(const ExecSpace &space, const SView sv, RView rv, MPI_Comm comm) -> Req<MpiSpace> {
+  using ST = typename SView::non_const_value_type;
+  using RT = typename RView::non_const_value_type;
+  static_assert(std::is_same_v<ST, RT>, "KokkosComm::mpi::iallgather: View value types must be identical");
+  Kokkos::Tools::pushRegion("KokkosComm::mpi::iallgather");
+
+  fail_if(!is_contiguous(sv) || !is_contiguous(rv),
+          "KokkosComm::mpi::iallgather: unimplemented for non-contiguous views");
+
+  // Sync: Work in space may have been used to produce view data.
+  space.fence("fence before non-blocking all-gather");
+
+  Req<MpiSpace> req;
+  // All ranks send/recv same count
+  MPI_Iallgather(data_handle(sv), span(sv), Impl::mpi_type_v<ST>, data_handle(rv), span(sv), Impl::mpi_type_v<RT>, comm,
+                 &req.mpi_request());
+  req.extend_view_lifetime(sv);
+  req.extend_view_lifetime(rv);
+
+  Kokkos::Tools::popRegion();
+  return req;
+}
+
 template <KokkosView SendView, KokkosView RecvView>
 void allgather(const SendView &sv, const RecvView &rv, MPI_Comm comm) {
   Kokkos::Tools::pushRegion("KokkosComm::Mpi::allgather");