Add NCCL backend

.cscs-ci/default.yaml

@@ -4,7 +4,7 @@ include:
 variables:
   BASE_IMAGE: jfrog.svc.cscs.ch/docker-group-csstaff/alps-images/ngc-pytorch:26.01-py3-alps4-dev
   SPACK_SHA: v1.1.1
-  SPACK_PACKAGES_SHA: bc93746ce936d6653271b6e98f6df6ee28f64e84 # develop on 2026-03-25
+  SPACK_PACKAGES_SHA: 5f24787b5cd3c2356d9a8188b989ceb5307299c6 # https://github.com/msimberg/spack-packages/tree/oomph-nccl
   FF_TIMESTAMPS: true
 
 .build_deps_template:
@@ -25,13 +25,12 @@ variables:
     reports:
       dotenv: base-${BACKEND}.env
 
-# TODO: NCCL will be enabled in https://github.com/ghex-org/oomph/pull/55
-# build_deps_nccl:
-#   variables:
-#     BACKEND: nccl
-#   extends:
-#     - .container-builder-cscs-gh200
-#     - .build_deps_template
+build_deps_nccl:
+  variables:
+    BACKEND: nccl
+  extends:
+    - .container-builder-cscs-gh200
+    - .build_deps_template
 
 build_deps_mpi:
   variables:
@@ -69,14 +68,13 @@ build_deps_ucx:
     reports:
       dotenv: build-${BACKEND}.env
 
-# TODO: NCCL will be enabled in https://github.com/ghex-org/oomph/pull/55
-# build_nccl:
-#   variables:
-#     BACKEND: nccl
-#   extends: .build_template
-#   needs:
-#     - job: build_deps_nccl
-#       artifacts: true
+build_nccl:
+  variables:
+    BACKEND: nccl
+  extends: .build_template
+  needs:
+    - job: build_deps_nccl
+      artifacts: true
 
 build_mpi:
   variables:
@@ -136,20 +134,19 @@ build_ucx:
     - sleep 1
     - ctest --test-dir /oomph/build -L "parallel-ranks-4" --output-on-failure --timeout 60
 
-# TODO: NCCL will be enabled in https://github.com/ghex-org/oomph/pull/55
-# test_serial_nccl:
-#   extends: .test_serial_template
-#   needs:
-#     - job: build_nccl
-#       artifacts: true
-#   image: $BUILD_IMAGE
+test_serial_nccl:
+  extends: .test_serial_template
+  needs:
+    - job: build_nccl
+      artifacts: true
+  image: $BUILD_IMAGE
 
-# test_parallel_nccl:
-#   extends: .test_parallel_template
-#   needs:
-#     - job: build_nccl
-#       artifacts: true
-#   image: $BUILD_IMAGE
+test_parallel_nccl:
+  extends: .test_parallel_template
+  needs:
+    - job: build_nccl
+      artifacts: true
+  image: $BUILD_IMAGE
 
 test_serial_mpi:
   extends: .test_serial_template

CMakeLists.txt

@@ -80,6 +80,11 @@ include(oomph_ucx)
 # ---------------------------------------------------------------------
 include(oomph_libfabric)
 
+# ---------------------------------------------------------------------
+# oomph NCCL variant
+# ---------------------------------------------------------------------
+include(oomph_nccl)
+
 # ---------------------------------------------------------------------
 # main src subdir
 # ---------------------------------------------------------------------
@@ -142,6 +147,7 @@ install(
         ${CMAKE_CURRENT_BINARY_DIR}/oomphConfig.cmake
         ${CMAKE_CURRENT_BINARY_DIR}/oomphConfigVersion.cmake
         ${CMAKE_CURRENT_LIST_DIR}/cmake/FindLibfabric.cmake
+        ${CMAKE_CURRENT_LIST_DIR}/cmake/FindNCCL.cmake
         ${CMAKE_CURRENT_LIST_DIR}/cmake/FindUCX.cmake
         ${CMAKE_CURRENT_LIST_DIR}/cmake/FindPMIx.cmake
     DESTINATION

README.md

@@ -4,9 +4,9 @@
 # OOMPH
 
 **Oomph** is a library for enabling high performance point-to-point, asynchronous communication over
-different fabrics.  It leverages the ubiquitos MPI library as well as UCX and Libfabric.  Both
-device and host memory are supported. Under the hood it uses
-[hwmalloc](https://github.com/ghex-org/hwmalloc) for memory registration.
+different fabrics. It leverages the ubiquitous MPI library as well as UCX and Libfabric. Both
+device and host memory are supported. A subset of functionality is also supported with NCCL. Under
+the hood it uses [hwmalloc](https://github.com/ghex-org/hwmalloc) for memory registration.
 
 **selling points**
 - lightweight, fast
@@ -136,6 +136,68 @@ comm.progress();
 // or progress until some event is triggered
 while(!completed) { comm.progress(); }
 ```
+
+### Groups
+
+Communicators expose group functionality as provided by NCCL (with
+ncclGroupStart and ncclGroupEnd). For non-NCCL backends the group functionality
+is a no-op. For NCCL using the group functionality can be a both a requirement
+to avoid deadlocks (communication within a group can make progress
+independently, while outside of a group communication is ordered) and for
+performance (a single device kernel is submitted for a NCCL group). 
+
+Groups are created by explicitly starting and ending the group:
+
+```cpp
+comm.start_group();
+oomph::send_request sreq = comm.send(smsg, 1, 0);
+oomph::recv_request rreq = comm.recv(rmsg, 1, 0);
+comm.end_group();
+
+// With NCCL, no progress will be made until after the group ends
+sreq.wait();
+rreq.wait();
+```
+
+### Stream awareness
+
+Some backend implementations can schedule communication on a GPU stream.
+Currently only the NCCL backend makes use of this. All other backends ignore
+the stream argument. To query if a backend is stream-aware use the
+`is_stream_aware` member query on a communicator. The stream can be passed as
+an optional last parameter to `send` or `recv`:
+
+```cpp
+if (comm.is_stream_aware()) {
+    # Schedule communication on the default CUDA stream if the backend is
+    # stream aware
+    cudaStream_t stream = 0;
+    oomph::send_request comm.send(msg, 1, 0, stream);
+}
+```
+
+### NCCL restrictions
+
+NCCL has significantly different semantics from MPI, libfabric, and UCX which
+is reflected in a number of restrictions on how the NCCL communicator can be
+used:
+
+- Tags are not supported by NCCL and ignored by the backend. Communication
+  order on different ranks must match (except within NCCL groups where there is
+  some flexibility). This also means that e.g. recv should not be called before
+  send unless within a NCCL group.
+- The `thread_safe` option for the NCCL communicator is not supported because
+  of the above ordering restrictions.
+- Cancellation is not supported.
+- `wait` and `progress` are disallowed when a NCCL group is active as no
+  progress can be made until a NCCL group is ended and submitted.
+- Send/recv to own rank is supported within NCCL groups. Outside of a group,
+  self-send/recv will throw an exception because NCCL's order-based matching
+  requires the atomic submission that groups provide.
+
+The NCCL backend is primarily designed for use in GHEX where these differences
+can be hidden from the user.
+
 ## Acknowledgments
 This work was financially supported by the PRACE project funded in part by the EU's Horizon 2020
 Research and Innovation programme (2014-2020) under grant agreement 823767.

cmake/FindNCCL.cmake

@@ -0,0 +1,78 @@
+# From https://github.com/pytorch/gloo/blob/main/cmake/Modules/Findnccl.cmake.
+
+# Try to find NCCL
+#
+# The following variables are optionally searched for defaults
+#  NCCL_ROOT_DIR: 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_ROOT_DIR $ENV{NCCL_ROOT_DIR} CACHE PATH "Folder contains NVIDIA NCCL")
+
+find_path(NCCL_INCLUDE_DIR
+  NAMES nccl.h
+  HINTS
+  ${NCCL_INCLUDE_DIR}
+  ${NCCL_ROOT_DIR}
+  ${NCCL_ROOT_DIR}/include
+  ${CUDA_TOOLKIT_ROOT_DIR}/include)
+
+if(DEFINED ENV{USE_STATIC_NCCL} AND NOT "$ENV{USE_STATIC_NCCL}" STREQUAL "")
+  message(STATUS "USE_STATIC_NCCL detected. Linking against static NCCL library")
+  set(_use_static_nccl ON)
+  set(NCCL_LIBNAME "libnccl_static.a")
+else()
+  set(_use_static_nccl OFF)
+  set(NCCL_LIBNAME "nccl")
+endif()
+
+find_library(NCCL_LIBRARY
+  NAMES ${NCCL_LIBNAME}
+  HINTS
+  ${NCCL_LIB_DIR}
+  ${NCCL_ROOT_DIR}
+  ${NCCL_ROOT_DIR}/lib
+  ${NCCL_ROOT_DIR}/lib/x86_64-linux-gnu
+  ${NCCL_ROOT_DIR}/lib64
+  ${CUDA_TOOLKIT_ROOT_DIR}/lib64)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIR NCCL_LIBRARY)
+
+if (NCCL_FOUND)
+  set(NCCL_HEADER_FILE "${NCCL_INCLUDE_DIR}/nccl.h")
+  message(STATUS "Determining NCCL version from the header file: ${NCCL_HEADER_FILE}")
+  file (STRINGS ${NCCL_HEADER_FILE} NCCL_MAJOR_VERSION_DEFINED
+        REGEX "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+[0-9]+.*$" LIMIT_COUNT 1)
+  if (NCCL_MAJOR_VERSION_DEFINED)
+    string (REGEX REPLACE "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+" ""
+            NCCL_MAJOR_VERSION ${NCCL_MAJOR_VERSION_DEFINED})
+    message(STATUS "NCCL_MAJOR_VERSION: ${NCCL_MAJOR_VERSION}")
+  endif()
+  set(NCCL_INCLUDE_DIRS ${NCCL_INCLUDE_DIR})
+  set(NCCL_LIBRARIES ${NCCL_LIBRARY})
+  message(STATUS "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})")
+  mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES)
+
+  if(NOT TARGET NCCL::nccl AND NCCL_FOUND)
+    if(_use_static_nccl)
+      add_library(NCCL::nccl STATIC IMPORTED)
+    else()
+      add_library(NCCL::nccl SHARED IMPORTED)
+    endif()
+    set_target_properties(NCCL::nccl PROPERTIES
+      IMPORTED_LOCATION ${NCCL_LIBRARIES}
+      INTERFACE_INCLUDE_DIRECTORIES ${NCCL_INCLUDE_DIRS}
+    )
+  endif()
+endif()
+

cmake/oomphConfig.cmake.in

@@ -19,5 +19,8 @@ if (@OOMPH_WITH_LIBFABRIC@)
     #set(LIBFABRIC_INCLUDE_DIR @ULIBFABRIC_INCLUDE_DIRS@)
     find_dependency(Libfabric)
 endif()
+if (@OOMPH_WITH_NCCL@)
+    find_dependency(NCCL)
+endif()
 include(${CMAKE_CURRENT_LIST_DIR}/oomph-targets.cmake)
 

cmake/oomph_nccl.cmake

@@ -0,0 +1,19 @@
+# set all NCCL related options and values
+
+#------------------------------------------------------------------------------
+# Enable NCCL support
+#------------------------------------------------------------------------------
+set(OOMPH_WITH_NCCL OFF CACHE BOOL "Build with NCCL backend")
+
+if (OOMPH_WITH_NCCL)
+    find_package(CUDAToolkit REQUIRED)
+    find_package(NCCL REQUIRED)
+    add_library(oomph_nccl SHARED)
+    add_library(oomph::nccl ALIAS oomph_nccl)
+    oomph_shared_lib_options(oomph_nccl)
+    target_link_libraries(oomph_nccl PUBLIC NCCL::nccl CUDA::cudart)
+    install(TARGETS oomph_nccl
+        EXPORT oomph-targets
+        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR})
+endif()