Add support for NCCL in communication_object

.cscs-ci/container/build.Containerfile

@@ -16,6 +16,7 @@ RUN spack -e ci build-env ghex -- \
             -DGHEX_USE_BUNDLED_GTEST=ON \
             -DGHEX_USE_BUNDLED_OOMPH=OFF \
             -DGHEX_USE_BUNDLED_GRIDTOOLS=OFF \
+            -DCMAKE_CUDA_ARCHITECTURES=90 \
             -DGHEX_USE_GPU=ON \
             -DGHEX_GPU_TYPE=NVIDIA \
             -DGHEX_BUILD_PYTHON_BINDINGS=ON \

.cscs-ci/default.yaml

@@ -2,9 +2,9 @@ include:
   - remote: 'https://gitlab.com/cscs-ci/recipes/-/raw/master/templates/v2/.ci-ext.yml'
 
 variables:
-  BASE_IMAGE: jfrog.svc.cscs.ch/docker-group-csstaff/alps-images/ngc-pytorch:26.01-py3-alps4-dev
+  BASE_IMAGE: jfrog.svc.cscs.ch/docker-group-csstaff/alps-images/ngc-pytorch:26.01-py3-alps3
   SPACK_SHA: v1.1.1
-  SPACK_PACKAGES_SHA: a010c65289743f900bdbbfb840e4d1876c24e93f # develop on 2025-05-08
+  SPACK_PACKAGES_SHA: 8ea120fe82c02737dddef32451edf88929f266ff # https://github.com/msimberg/spack-packages/tree/oomph-nccl
   FF_TIMESTAMPS: true
 
 .build_deps_template:
@@ -25,6 +25,12 @@ variables:
     reports:
       dotenv: base-${BACKEND}.env
 
+build_deps_nccl:
+  variables:
+    BACKEND: nccl
+  extends:
+    - .build_deps_template
+
 build_deps_mpi:
   extends: .build_deps_template
   variables:
@@ -53,6 +59,14 @@ build_deps_libfabric:
     reports:
       dotenv: build-${BACKEND}.env
 
+build_nccl:
+  extends: .build_template
+  variables:
+    BACKEND: nccl
+  needs:
+    - job: build_deps_nccl
+      artifacts: true
+
 build_mpi:
   extends: .build_template
   variables:
@@ -81,7 +95,7 @@ build_libfabric:
   extends: .container-runner-clariden-gh200
   variables:
     SLURM_GPUS_PER_TASK: 1
-    SLURM_TIMELIMIT: '5:00'
+    SLURM_TIMELIMIT: '15:00'
     SLURM_PARTITION: normal
     SLURM_MPI_TYPE: pmix
     SLURM_NETWORK: disable_rdzv_get
@@ -90,13 +104,15 @@ build_libfabric:
     PMIX_MCA_psec: native
     PMIX_MCA_gds: "^shmem2"
     USE_MPI: NO
+    NCCL_DEBUG: trace
+    NCCL_PXN_DISABLE: 1
 
 .test_serial_template:
   extends: .test_template_base
   variables:
     SLURM_NTASKS: 1
   script:
-    - spack -e ci build-env ghex -- ctest --test-dir /ghex/build -L "serial" --output-on-failure --timeout 60 --parallel 8
+    - spack -e ci build-env ghex -- ctest --test-dir /ghex/build -L "serial" --verbose --output-on-failure --timeout 600 --parallel 8
 
 .test_parallel_template:
   extends: .test_template_base
@@ -105,12 +121,18 @@ build_libfabric:
     # writing inside the container.
     - if [[ "${SLURM_LOCALID}" == 0 ]]; then rm -rf /ghex/build/Testing; mkdir /tmp/Testing; ln -s /tmp/Testing /ghex/build/Testing; fi
     - until [[ -L /ghex/build/Testing ]]; do sleep 1; done
-    - spack -e ci build-env ghex -- ctest --test-dir /ghex/build -L "parallel-ranks-${SLURM_NTASKS}" --output-on-failure --timeout 60
+    - spack -e ci build-env ghex -- ctest --test-dir /ghex/build -L "parallel-ranks-${SLURM_NTASKS}" --verbose --output-on-failure --timeout 600
 
 .test_parallel_job:
   extends: .test_parallel_template
   image: $BUILD_IMAGE
 
+.test_parallel_nccl:
+  extends: .test_parallel_job
+  needs:
+    - job: build_nccl
+      artifacts: true
+
 .test_parallel_mpi:
   extends: .test_parallel_job
   needs:
@@ -129,6 +151,28 @@ build_libfabric:
 #    - job: build_libfabric
 #      artifacts: true
 
+test_serial_nccl:
+  extends: .test_serial_template
+  needs:
+    - job: build_nccl
+      artifacts: true
+  image: $BUILD_IMAGE
+
+test_parallel_2_nccl:
+  extends: .test_parallel_nccl
+  variables:
+    SLURM_NTASKS: 2
+
+test_parallel_4_nccl:
+  extends: .test_parallel_nccl
+  variables:
+    SLURM_NTASKS: 4
+
+test_parallel_6_nccl:
+  extends: .test_parallel_nccl
+  variables:
+    SLURM_NTASKS: 6
+
 test_serial_mpi:
   extends: .test_serial_template
   needs:

.cscs-ci/spack/libfabric.yaml

@@ -4,3 +4,8 @@ spack:
   view: false
   concretizer:
     unify: true
+  packages:
+    oomph:
+      require: '@git.365fab9af9538694668fc8516750bfb0e96b6be2=main'
+      package_attributes:
+        git: "https://github.com/msimberg/oomph.git"

.cscs-ci/spack/mpi.yaml

@@ -4,3 +4,8 @@ spack:
   view: false
   concretizer:
     unify: true
+  packages:
+    oomph:
+      require: '@git.365fab9af9538694668fc8516750bfb0e96b6be2=main'
+      package_attributes:
+        git: "https://github.com/msimberg/oomph.git"

.cscs-ci/spack/nccl.yaml

@@ -0,0 +1,11 @@
+spack:
+  specs:
+  - ghex@master backend=nccl +cuda cuda_arch=90a +python
+  view: false
+  concretizer:
+    unify: true
+  packages:
+    oomph:
+      require: '@git.3d7f65888e10ba1689257184005b12e53c907738=main'
+      package_attributes:
+        git: "https://github.com/msimberg/oomph.git"

.cscs-ci/spack/ucx.yaml

@@ -4,3 +4,8 @@ spack:
   view: false
   concretizer:
     unify: true
+  packages:
+    oomph:
+      require: '@git.365fab9af9538694668fc8516750bfb0e96b6be2=main'
+      package_attributes:
+        git: "https://github.com/msimberg/oomph.git"

.gitignore

@@ -48,3 +48,4 @@ doc_src/_build/
 __pycache__
 _build
 .venv*/
+_nix_build/

CMakeLists.txt

@@ -172,6 +172,8 @@ if(GHEX_USE_BUNDLED_OOMPH)
         set_target_properties(oomph_libfabric PROPERTIES INSTALL_RPATH "${rpath_origin}")
     elseif (GHEX_TRANSPORT_BACKEND STREQUAL "UCX")
         set_target_properties(oomph_ucx PROPERTIES INSTALL_RPATH "${rpath_origin}")
+    elseif (GHEX_TRANSPORT_BACKEND STREQUAL "NCCL")
+        set_target_properties(oomph_nccl PROPERTIES INSTALL_RPATH "${rpath_origin}")
     else()
         set_target_properties(oomph_mpi PROPERTIES INSTALL_RPATH "${rpath_origin}")
     endif()

README.md

@@ -58,7 +58,7 @@ make test
 | `GHEX_PYTHON_LIB_PATH=`           | `<path>`                | `${CMAKE_INSTALL_PREFIX}/<python-site-packages-path>` | Installation directory for GHEX's Python package
 | `GHEX_WITH_TESTING=`              | `{ON, OFF}`             | `OFF`                                                 | Build unit tests
 | `GHEX_USE_XPMEM=`                 | `{ON, OFF}`             | `OFF`                                                 | Use Xpmem
-| `GHEX_TRANSPORT_BACKEND=`         | `{MPI, UCX, LIBFABRIC}` | `MPI`                                                 | Choose transport backend
+| `GHEX_TRANSPORT_BACKEND=`         | `{MPI, UCX, LIBFABRIC, NCCL}` | `MPI`                                                 | Choose transport backend
 
 ### Pip Install
 
@@ -75,7 +75,7 @@ python -m pip install ghex
 | `GHEX_USE_GPU=`           | `{ON, OFF}`             | `OFF`                              | Enable GPU
 | `GHEX_GPU_TYPE=`          | `{AUTO, NVIDIA, AMD}`   | `AUTO`                             | Choose GPU type
 | `GHEX_GPU_ARCH=`          | list of archs           | `"60;70;75;80"`/ `"gfx900;gfx906"` | GPU architecture
-| `GHEX_TRANSPORT_BACKEND=` | `{MPI, UCX, LIBFABRIC}` | `MPI`                              | Choose transport backend
+| `GHEX_TRANSPORT_BACKEND=` | `{MPI, UCX, LIBFABRIC, NCCL}` | `MPI`                              | Choose transport backend
 
 
 ## Contributing

cmake/ghex_external_dependencies.cmake

@@ -43,8 +43,8 @@ endif()
 # ---------------------------------------------------------------------
 # oomph setup
 # ---------------------------------------------------------------------
-set(GHEX_TRANSPORT_BACKEND "MPI" CACHE STRING "Choose the backend type: MPI | UCX | LIBFABRIC")
-set_property(CACHE GHEX_TRANSPORT_BACKEND PROPERTY STRINGS "MPI" "UCX" "LIBFABRIC")
+set(GHEX_TRANSPORT_BACKEND "MPI" CACHE STRING "Choose the backend type: MPI | UCX | LIBFABRIC | NCCL")
+set_property(CACHE GHEX_TRANSPORT_BACKEND PROPERTY STRINGS "MPI" "UCX" "LIBFABRIC" "NCCL")
 cmake_dependent_option(GHEX_USE_BUNDLED_OOMPH "Use bundled oomph." ON "GHEX_USE_BUNDLED_LIBS" OFF)
 if(GHEX_USE_BUNDLED_OOMPH)
     set(OOMPH_GIT_SUBMODULE OFF CACHE BOOL "")
@@ -53,6 +53,11 @@ if(GHEX_USE_BUNDLED_OOMPH)
         set(OOMPH_WITH_LIBFABRIC ON CACHE BOOL "Build with LIBFABRIC backend")
     elseif(GHEX_TRANSPORT_BACKEND STREQUAL "UCX")
         set(OOMPH_WITH_UCX ON CACHE BOOL "Build with UCX backend")
+    elseif(GHEX_TRANSPORT_BACKEND STREQUAL "NCCL")
+        set(OOMPH_WITH_NCCL ON CACHE BOOL "Build with NCCL backend")
+        if(NOT GHEX_USE_GPU)
+          message(FATAL_ERROR "GHEX_TRANSPORT_BACKEND=NCCL requires GHEX_USE_GPU=ON but GHEX_USE_GPU=OFF")
+        endif()
     endif()
     if(GHEX_USE_GPU)
         set(HWMALLOC_ENABLE_DEVICE ON CACHE BOOL "True if GPU support shall be enabled")
@@ -70,6 +75,9 @@ if(GHEX_USE_BUNDLED_OOMPH)
     if(TARGET oomph_ucx)
         add_library(oomph::oomph_ucx ALIAS oomph_ucx)
     endif()
+    if(TARGET oomph_nccl)
+        add_library(oomph::oomph_nccl ALIAS oomph_nccl)
+    endif()
     if(TARGET oomph_libfabric)
         add_library(oomph::oomph_libfabric ALIAS oomph_libfabric)
     endif()
@@ -82,6 +90,8 @@ function(ghex_link_to_oomph target)
         target_link_libraries(${target} PRIVATE oomph::oomph_libfabric)
     elseif (GHEX_TRANSPORT_BACKEND STREQUAL "UCX")
         target_link_libraries(${target} PRIVATE oomph::oomph_ucx)
+    elseif (GHEX_TRANSPORT_BACKEND STREQUAL "NCCL")
+        target_link_libraries(${target} PRIVATE oomph::oomph_nccl)
     else()
         target_link_libraries(${target} PRIVATE oomph::oomph_mpi)
     endif()