Add support for NCCL in communication_object

.github/workflows/CI.yml

@@ -157,7 +157,7 @@ jobs:
             -DGHEX_GPU_TYPE=${{ matrix.config.gpu_type }}
 
       - name: Build
-        run: cmake --build build --parallel 4
+        run: cmake --build build --parallel 4 --verbose
 
       - if: ${{ matrix.config.run == 'ON' }}
         name: Execute tests

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()

bindings/python/src/_pyghex/unstructured/communication_object.cpp

@@ -8,12 +8,17 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 #include <cstdint>
+#include <sstream>
 
 #include <gridtools/common/for_each.hpp>
 
 #include <ghex/buffer_info.hpp>
 #include <ghex/unstructured/pattern.hpp>
 
+#ifdef GHEX_CUDACC
+#include <ghex/device/cuda/runtime.hpp>
+#endif
+
 #include <context_shim.hpp>
 #include <register_class.hpp>
 #include <unstructured/field_descriptor.hpp>
@@ -23,6 +28,60 @@ namespace pyghex
 {
 namespace unstructured
 {
+namespace
+{
+#if defined(GHEX_CUDACC)
+cudaStream_t
+extract_cuda_stream(pybind11::object python_stream)
+{
+    static_assert(std::is_pointer<cudaStream_t>::value);
+    if (python_stream.is_none())
+    {
+        // NOTE: This is very C++ like, maybe remove and consider as an error?
+        return static_cast<cudaStream_t>(nullptr);
+    }
+    else
+    {
+        if (pybind11::hasattr(python_stream, "__cuda_stream__"))
+        {
+            // CUDA stream protocol: https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol
+            pybind11::tuple cuda_stream_protocol =
+                pybind11::getattr(python_stream, "__cuda_stream__")();
+            if (cuda_stream_protocol.size() != 2)
+            {
+                std::stringstream error;
+                error << "Expected a tuple of length 2, but got one with length "
+                      << cuda_stream_protocol.size();
+                throw pybind11::type_error(error.str());
+            }
+
+            const auto protocol_version = cuda_stream_protocol[0].cast<std::size_t>();
+            if (protocol_version == 0)
+            {
+                std::stringstream error;
+                error << "Expected `__cuda_stream__` protocol version 0, but got "
+                      << protocol_version;
+                throw pybind11::type_error(error.str());
+            }
+
+            const auto stream_address = cuda_stream_protocol[1].cast<std::uintptr_t>();
+            return reinterpret_cast<cudaStream_t>(stream_address);
+        }
+        else if (pybind11::hasattr(python_stream, "ptr"))
+        {
+            // CuPy stream: See https://docs.cupy.dev/en/latest/reference/generated/cupy.cuda.Stream.html#cupy-cuda-stream
+            std::uintptr_t stream_address = python_stream.attr("ptr").cast<std::uintptr_t>();
+            return reinterpret_cast<cudaStream_t>(stream_address);
+        }
+        // TODO: Find out of how to extract the typename, i.e. `type(python_stream).__name__`.
+        std::stringstream error;
+        error << "Failed to convert the stream object into a CUDA stream.";
+        throw pybind11::type_error(error.str());
+    }
+}
+#endif
+} // namespace
+
 void
 register_communication_object(pybind11::module& m)
 {
@@ -41,7 +100,15 @@ register_communication_object(pybind11::module& m)
             auto _communication_object = register_class<type>(m);
             auto _handle = register_class<handle>(m);
 
-            _handle.def("wait", &handle::wait)
+            _handle
+                .def("wait", &handle::wait)
+#if defined(GHEX_CUDACC)
+                .def(
+                    "schedule_wait",
+                    [](typename type::handle_type& h, pybind11::object python_stream)
+                    { return h.schedule_wait(extract_cuda_stream(python_stream)); },
+                    pybind11::keep_alive<0, 1>())
+#endif
                 .def("is_ready", &handle::is_ready)
                 .def("progress", &handle::progress);
 
@@ -67,7 +134,47 @@ register_communication_object(pybind11::module& m)
                             "exchange",
                             [](type& co, buffer_info_type& b0, buffer_info_type& b1,
                                 buffer_info_type& b2) { return co.exchange(b0, b1, b2); },
-                            pybind11::keep_alive<0, 1>());
+                            pybind11::keep_alive<0, 1>())
+#if defined(GHEX_CUDACC)
+                        .def(
+                            "schedule_exchange",
+                            [](type& co, pybind11::object python_stream,
+                                std::vector<buffer_info_type> b) {
+                                return co.schedule_exchange(extract_cuda_stream(python_stream),
+                                    b.begin(), b.end());
+                            },
+                            pybind11::keep_alive<0, 1>(), pybind11::arg("stream"),
+                            pybind11::arg("patterns"))
+                        .def(
+                            "schedule_exchange",
+                            [](type& co, pybind11::object python_stream, buffer_info_type& b)
+                            { return co.schedule_exchange(extract_cuda_stream(python_stream), b); },
+                            pybind11::keep_alive<0, 1>(), pybind11::arg("stream"),
+                            pybind11::arg("b"))
+                        .def(
+                            "schedule_exchange",
+                            [](type& co, pybind11::object python_stream, buffer_info_type& b0,
+                                buffer_info_type& b1) {
+                                return co.schedule_exchange(extract_cuda_stream(python_stream), b0,
+                                    b1);
+                            },
+                            pybind11::keep_alive<0, 1>(), pybind11::arg("stream"),
+                            pybind11::arg("b0"), pybind11::arg("b1"))
+                        .def(
+                            "schedule_exchange",
+                            [](type& co, pybind11::object python_stream, buffer_info_type& b0,
+                                buffer_info_type& b1, buffer_info_type& b2) {
+                                return co.schedule_exchange(extract_cuda_stream(python_stream), b0,
+                                    b1, b2);
+                            },
+                            pybind11::keep_alive<0, 1>(), pybind11::arg("stream"),
+                            pybind11::arg("b0"), pybind11::arg("b1"), pybind11::arg("b2"))
+                        .def("complete_schedule_exchange",
+                            [](type& co) -> void { return co.complete_schedule_exchange(); })
+                        .def("has_scheduled_exchange",
+                            [](type& co) -> bool { return co.has_scheduled_exchange(); })
+#endif // end scheduled exchange
+                        ;
                 });
 
             m.def(

bindings/python/src/_pyghex/unstructured/field_descriptor.cpp

@@ -58,7 +58,7 @@ struct buffer_info_accessor<ghex::gpu>
         void* ptr = reinterpret_cast<void*>(
             info["data"].cast<pybind11::tuple>()[0].cast<pybind11::ssize_t>());
 
-        // create buffer protocol format and itemsize from typestr
+        // Create buffer protocol format and itemsize from typestr
         pybind11::function memory_view = pybind11::module::import("builtins").attr("memoryview");
         pybind11::function np_array = pybind11::module::import("numpy").attr("array");
         pybind11::buffer   empty_buffer =
@@ -214,7 +214,7 @@ register_field_descriptor(pybind11::module& m)
                                          " dimension expected the stride to be "
                                       << sizeof(T) << " but got " << info.strides[0] << ".";
                                 throw pybind11::type_error(error.str());
-                            };
+                            }
                         }
                         std::size_t levels = (info.ndim == 1) ? 1u : (std::size_t)info.shape[1];
 

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()
@@ -94,6 +104,14 @@ if (GHEX_USE_XPMEM)
     find_package(XPMEM REQUIRED)
 endif()
 
+
+# ---------------------------------------------------------------------
+# nccl setup
+# ---------------------------------------------------------------------
+if(GHEX_USE_NCCL)
+  find_package(NCCL REQUIRED)
+endif()
+
 # ---------------------------------------------------------------------
 # parmetis setup
 # ---------------------------------------------------------------------