Feature: Add basic NCCL communication space backend

.gitignore

@@ -5,4 +5,6 @@ Testing
 kokkos
 docs/_build
 .python-version
-.venv
+.venv
+.cache
+compile_commands.json

CMakeLists.txt

@@ -22,8 +22,7 @@ cmake_minimum_required(VERSION 3.23)
 
 project(
   KokkosComm
-  LANGUAGES
-    CXX
+  LANGUAGES CXX
   VERSION 0.2.0
   DESCRIPTION "Experimental MPI interfaces (and more!) for the Kokkos C++ Performance Portability Programming ecosystem"
   HOMEPAGE_URL "https://kokkos.org/kokkos-comm/"
@@ -32,16 +31,21 @@ project(
 option(KokkosComm_ENABLE_PERFTESTS "Build KokkosComm perf tests" OFF)
 option(KokkosComm_ENABLE_TESTS "Build KokkosComm tests" OFF)
 option(KokkosComm_ENABLE_MPI "Build KokkosComm with MPI transport" ON)
+option(KokkosComm_ENABLE_NCCL "Build KokkosComm with NCCL transport" OFF)
 
 # 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)
+set(KOKKOSCOMM_ENABLE_NCCL ${KokkosComm_ENABLE_NCCL} CACHE BOOL "" FORCE)
 
 find_package(Kokkos REQUIRED)
 if(KOKKOSCOMM_ENABLE_MPI)
   find_package(MPI REQUIRED)
 endif()
+if(KOKKOSCOMM_ENABLE_NCCL)
+  find_package(NCCL REQUIRED)
+endif()
 
 add_subdirectory(src)
 if(KOKKOSCOMM_ENABLE_TESTS)

cmake/FindNCCL.cmake

@@ -0,0 +1,99 @@
+# Find the NCCL libraries.
+# Copied from pytorch: https://github.com/pytorch/pytorch/blob/main/cmake/Modules/FindNCCL.cmake
+#
+# The following variables are optionally searched for defaults
+#  NCCL_ROOT: Base directory where all NCCL components are found
+#  NCCL_INCLUDE_DIR: Directory where NCCL header is found
+#  NCCL_LIB_DIR: Directory where NCCL library is found
+#
+# The following are set after configuration is done:
+#  NCCL_FOUND
+#  NCCL_INCLUDE_DIRS
+#  NCCL_LIBRARIES
+#
+# The path hints include CUDA_TOOLKIT_ROOT_DIR seeing as some folks
+# install NCCL in the same location as the CUDA toolkit.
+# See https://github.com/caffe2/caffe2/issues/1601
+
+set(NCCL_INCLUDE_DIR $ENV{NCCL_INCLUDE_DIR} CACHE PATH "Folder contains NVIDIA NCCL headers")
+set(NCCL_LIB_DIR $ENV{NCCL_LIB_DIR} CACHE PATH "Folder contains NVIDIA NCCL libraries")
+set(NCCL_VERSION $ENV{NCCL_VERSION} CACHE STRING "Version of NCCL to build with")
+
+if($ENV{NCCL_ROOT_DIR})
+  message(WARNING "NCCL_ROOT_DIR is deprecated. Please set NCCL_ROOT instead.")
+endif()
+list(APPEND NCCL_ROOT $ENV{NCCL_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR})
+# Compatible layer for CMake <3.12. NCCL_ROOT will be accounted in for searching paths and libraries for CMake >=3.12.
+list(APPEND CMAKE_PREFIX_PATH ${NCCL_ROOT})
+
+find_path(NCCL_INCLUDE_DIRS NAMES nccl.h HINTS ${NCCL_INCLUDE_DIR})
+
+if(USE_STATIC_NCCL)
+  message(STATUS "USE_STATIC_NCCL is set. Linking with static NCCL library.")
+  set(NCCL_LIBNAME "nccl_static")
+  if(NCCL_VERSION) # Prefer the versioned library if a specific NCCL version is specified
+    set(CMAKE_FIND_LIBRARY_SUFFIXES ".a.${NCCL_VERSION}" ${CMAKE_FIND_LIBRARY_SUFFIXES})
+  endif()
+else()
+  set(NCCL_LIBNAME "nccl")
+  if(NCCL_VERSION) # Prefer the versioned library if a specific NCCL version is specified
+    set(CMAKE_FIND_LIBRARY_SUFFIXES ".so.${NCCL_VERSION}" ${CMAKE_FIND_LIBRARY_SUFFIXES})
+  endif()
+endif()
+
+find_library(NCCL_LIBRARIES NAMES ${NCCL_LIBNAME} HINTS ${NCCL_LIB_DIR})
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIRS NCCL_LIBRARIES)
+
+if(NCCL_FOUND) # obtaining NCCL version and some sanity checks
+  set(NCCL_HEADER_FILE "${NCCL_INCLUDE_DIRS}/nccl.h")
+  message(STATUS "Determining NCCL version from ${NCCL_HEADER_FILE}...")
+  set(OLD_CMAKE_REQUIRED_INCLUDES ${CMAKE_REQUIRED_INCLUDES})
+  list(APPEND CMAKE_REQUIRED_INCLUDES ${NCCL_INCLUDE_DIRS})
+  include(CheckCXXSymbolExists)
+  check_cxx_symbol_exists(NCCL_VERSION_CODE nccl.h NCCL_VERSION_DEFINED)
+
+  if(NCCL_VERSION_DEFINED)
+    set(file "${PROJECT_BINARY_DIR}/detect_nccl_version.cc")
+    file(
+      WRITE
+      ${file}
+      "
+      #include <iostream>
+      #include <nccl.h>
+      int main()
+      {
+        std::cout << NCCL_MAJOR << '.' << NCCL_MINOR << '.' << NCCL_PATCH << std::endl;
+
+        int x;
+        ncclGetVersion(&x);
+        return x == NCCL_VERSION_CODE;
+      }
+"
+    )
+    try_run(
+      NCCL_VERSION_MATCHED
+      compile_result
+      ${PROJECT_BINARY_DIR}
+      ${file}
+      RUN_OUTPUT_VARIABLE NCCL_VERSION_FROM_HEADER
+      CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${NCCL_INCLUDE_DIRS}"
+      LINK_LIBRARIES ${NCCL_LIBRARIES}
+    )
+    if(NOT NCCL_VERSION_MATCHED)
+      message(
+        FATAL_ERROR
+        "Found NCCL header version and library version do not match! \
+(include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES}) Please set NCCL_INCLUDE_DIR and NCCL_LIB_DIR manually."
+      )
+    endif()
+    message(STATUS "NCCL version: ${NCCL_VERSION_FROM_HEADER}")
+  else()
+    message(STATUS "NCCL version < 2.3.5-5")
+  endif()
+  set(CMAKE_REQUIRED_INCLUDES ${OLD_CMAKE_REQUIRED_INCLUDES})
+
+  message(STATUS "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})")
+  mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES)
+endif()

src/KokkosComm/CMakeLists.txt

@@ -14,8 +14,6 @@
 #
 #@HEADER
 
-set(TARGET_LIBRARY kokkoscomm)
-
 add_library(KokkosComm INTERFACE)
 add_library(KokkosComm::KokkosComm ALIAS KokkosComm)
 
@@ -26,7 +24,7 @@ target_sources(
     FILE_SET kokkoscomm_public_headers
     TYPE HEADERS
     BASE_DIRS ${PROJECT_SOURCE_DIR}/src
-    FILES KokkosComm.hpp collective.hpp concepts.hpp fwd.hpp point_to_point.hpp traits.hpp
+    FILES KokkosComm.hpp collective.hpp concepts.hpp fwd.hpp point_to_point.hpp traits.hpp reduction_op.hpp
 )
 
 # Implementation detail headers
@@ -36,7 +34,7 @@ target_sources(
     FILE_SET kokkoscomm_impl_headers
     TYPE HEADERS
     BASE_DIRS ${PROJECT_SOURCE_DIR}/src
-    FILES impl/KokkosComm_contiguous.hpp
+    FILES impl/KokkosComm_contiguous.hpp impl/send.hpp impl/recv.hpp impl/allgather.hpp impl/reduce.hpp
 )
 
 # Configuration header
@@ -83,14 +81,39 @@ if(KOKKOSCOMM_ENABLE_MPI)
   )
 endif()
 
+if(KOKKOSCOMM_ENABLE_NCCL)
+  # Public NCCL headers
+  target_sources(
+    KokkosComm
+    INTERFACE
+      FILE_SET kokkoscomm_nccl_headers
+      TYPE HEADERS
+      BASE_DIRS ${PROJECT_SOURCE_DIR}/src
+      FILES nccl/nccl.hpp nccl/handle.hpp nccl/req.hpp nccl/send.hpp nccl/recv.hpp nccl/allgather.hpp nccl/reduce.hpp
+  )
+
+  # Implementation detail NCCL headers
+  target_sources(
+    KokkosComm
+    INTERFACE
+      FILE_SET kokkoscomm_nccl_impl_headers
+      TYPE HEADERS
+      BASE_DIRS ${PROJECT_SOURCE_DIR}/src
+      FILES nccl/impl/pack_traits.hpp nccl/impl/packer.hpp nccl/impl/types.hpp
+  )
+endif()
+
 # --- COMPILE FLAGS --- #
 include(CheckCXXCompilerFlag)
 
 macro(kokkoscomm_check_and_add_compile_options)
   set(target ${ARGV0})
   set(flag ${ARGV1})
 
-  check_cxx_compiler_flag(${flag} HAS_${flag})
+  check_cxx_compiler_flag(
+    ${flag}
+    HAS_${flag}
+  )
   if(HAS_${flag})
     target_compile_options(${target} INTERFACE ${flag})
   endif()
@@ -100,7 +123,12 @@ endmacro()
 add_library(KokkosCommFlags INTERFACE)
 add_library(KokkosComm::KokkosCommFlags ALIAS KokkosCommFlags)
 target_compile_features(KokkosCommFlags INTERFACE cxx_std_20)
-set_target_properties(KokkosCommFlags PROPERTIES CXX_EXTENSIONS OFF)
+set_target_properties(
+  KokkosCommFlags
+  PROPERTIES
+    CXX_EXTENSIONS
+      OFF
+)
 
 kokkoscomm_check_and_add_compile_options(KokkosCommFlags -Wall)
 kokkoscomm_check_and_add_compile_options(KokkosCommFlags -Wextra)
@@ -113,14 +141,30 @@ kokkoscomm_check_and_add_compile_options(KokkosCommFlags -Wmissing-include-dirs)
 kokkoscomm_check_and_add_compile_options(KokkosCommFlags -Wno-gnu-zero-variadic-macro-arguments)
 
 # Linking
-target_link_libraries(KokkosComm INTERFACE KokkosComm::KokkosCommFlags Kokkos::kokkos)
+target_link_libraries(
+  KokkosComm
+  INTERFACE
+    KokkosComm::KokkosCommFlags
+    Kokkos::kokkos
+)
 if(KOKKOSCOMM_ENABLE_MPI)
   target_link_libraries(KokkosComm INTERFACE MPI::MPI_CXX)
 endif()
+if(KOKKOSCOMM_ENABLE_NCCL)
+  # TODO: determine if this is the right way to link with NCCL
+  target_link_libraries(
+    KokkosComm
+    INTERFACE
+      MPI::MPI_CXX
+      ${NCCL_LIBRARIES}
+  )
+endif()
 
 # Install library
 install(
-  TARGETS KokkosComm KokkosCommFlags
+  TARGETS
+    KokkosComm
+    KokkosCommFlags
   EXPORT KokkosCommTargets
   FILE_SET
   kokkoscomm_public_headers
@@ -131,5 +175,9 @@ install(
   FILE_SET
   kokkoscomm_mpi_impl_headers
   FILE_SET
+  kokkoscomm_nccl_headers
+  FILE_SET
+  kokkoscomm_nccl_impl_headers
+  FILE_SET
   kokkoscomm_config_headers
 )

src/KokkosComm/collective.hpp

@@ -16,19 +16,35 @@
 
 #pragma once
 
-#include <utility>
+#include <KokkosComm/fwd.hpp>
 
-#include <Kokkos_Core.hpp>
+#include <Kokkos_Core_fwd.hpp>
 
-#include "fwd.hpp"
-#include "concepts.hpp"
+#include <utility>
 
 namespace KokkosComm {
 
 template <KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           CommunicationSpace CommSpace   = DefaultCommunicationSpace>
-void barrier(Handle<ExecSpace, CommSpace> &&h) {
+void barrier(Handle<ExecSpace, CommSpace>&& h) {
   Impl::Barrier<ExecSpace, CommSpace>{std::forward<Handle<ExecSpace, CommSpace>>(h)};
 }
 
+namespace Experimental {
+
+template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp,
+          KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          CommunicationSpace CommSpace   = DefaultCommunicationSpace>
+auto reduce(Handle<ExecSpace, CommSpace>& h, const SendView& sv, RecvView& rv, int root) -> Req<CommSpace> {
+  return Impl::Reduce<SendView, RecvView, RedOp, ExecSpace, CommSpace>(h, sv, rv, root);
+}
+
+template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          CommunicationSpace CommSpace = DefaultCommunicationSpace>
+auto allgather(Handle<ExecSpace, CommSpace>& h, const SendView& sv, RecvView& rv) -> Req<CommSpace> {
+  return Impl::AllGather<SendView, RecvView, ExecSpace, CommSpace>(h, sv, rv);
+}
+
+}  // namespace Experimental
+
 }  // namespace KokkosComm

src/KokkosComm/concepts.hpp

@@ -23,9 +23,14 @@
 namespace KokkosComm {
 
 namespace Impl {
+//
 // fallback - most types are not a KokkosComm transport
 template <typename T>
 struct is_communication_space : public std::false_type {};
+
+template <typename T>
+struct is_reduction_operator : public std::false_type {};
+
 }  // namespace Impl
 
 template <typename T>
@@ -37,4 +42,7 @@ concept KokkosExecutionSpace = Kokkos::is_execution_space_v<T>;
 template <typename T>
 concept CommunicationSpace = KokkosComm::Impl::is_communication_space<T>::value;
 
+template <typename T>
+concept ReductionOperator = KokkosComm::Impl::is_reduction_operator<T>::value;
+
 }  // namespace KokkosComm

src/KokkosComm/fwd.hpp

@@ -16,16 +16,21 @@
 
 #pragma once
 
-#include <vector>
-
-#include "concepts.hpp"
+#include <KokkosComm/concepts.hpp>
 #include <KokkosComm/config.hpp>
+#include <KokkosComm/reduction_op.hpp>
 
 namespace KokkosComm {
+
 #if defined(KOKKOSCOMM_ENABLE_MPI)
 class Mpi;
 using DefaultCommunicationSpace  = Mpi;
 using FallbackCommunicationSpace = Mpi;
+#elif defined(KOKKOSCOMM_ENABLE_NCCL)
+class Mpi;
+class Nccl;
+using DefaultCommunicationSpace  = Nccl;
+using FallbackCommunicationSpace = Mpi;
 #else
 #error at least one transport must be defined
 #endif
@@ -42,13 +47,29 @@ namespace Impl {
 template <KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           CommunicationSpace CommSpace = DefaultCommunicationSpace>
 struct Recv;
+
 template <KokkosView SendView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           CommunicationSpace CommSpace = DefaultCommunicationSpace>
 struct Send;
+
 template <KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
           CommunicationSpace CommSpace   = DefaultCommunicationSpace>
 struct Barrier;
 
 }  // namespace Impl
 
-}  // namespace KokkosComm
+// Allgather and Reduce are currently experimental functions
+namespace Experimental::Impl {
+
+template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          CommunicationSpace CommSpace = DefaultCommunicationSpace>
+struct AllGather;
+
+template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp,
+          KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
+          CommunicationSpace CommSpace   = DefaultCommunicationSpace>
+struct Reduce;
+
+}  // namespace Experimental::Impl
+
+}  // namespace KokkosComm