Add multiple GPUs support for the Python interface

.github/workflows/gh-pages.yml

@@ -16,6 +16,7 @@ jobs:
         shell: bash
         run: |
           module load python3/3.12.3 opensn/gcc/14 doxygen/1.13.2 pandoc/3.6.4
+          export CMAKE_ARGS="-DOPENSN_WITH_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=89 -DOPENSN_WITH_PYTHON_MODULE=ON"
           python3 setup.py build_ext --inplace
           cd doc
           make html

CMakeLists.txt

@@ -279,7 +279,6 @@ target_link_libraries(libopensn
     MPI::MPI_CXX
 )
 if (OPENSN_WITH_CUDA)
-  target_link_libraries(libopensncuda PRIVATE ${CUDA_LIBRARIES} CUDA::cublas)
   target_link_libraries(libopensn INTERFACE libopensncuda)
 endif()
 

doc/source/pyapi/index.rst

@@ -397,3 +397,26 @@ Argument vector
 
    context.InitializeWithArgv
    context.Finalize
+
+GPU configuration
+^^^^^^^^^^^^^^^^^
+
+.. important::
+
+   This API is only available when ``pyopensn.can_support_gpus`` is true.
+
+.. note::
+
+   The GPU assignment functionality is intended for use on single-user workstations.
+
+   It should not be used for production runs on large systems. When multiple GPUs are available on a
+   node, the mapping between MPI ranks and GPUs should be handled by the job scheduler.
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: noinit.rst
+
+   device.get_device_count
+   device.get_current_device
+   device.set_device

external/caribou/cuda/device.hpp

@@ -6,6 +6,7 @@
 #pragma once
 
 #include <cstdint>  // std::uint32_t
+#include <cstdio>   // std::printf
 
 #include "exception.hpp"  // cuda::check_cuda_error
 
@@ -21,4 +22,16 @@ inline std::uint32_t get_num_gpus(void) {
     return static_cast<std::uint32_t>(count);
 }
 
+/** @brief Get the current device number.*/
+inline int get_current_device(void) {
+    int device_num;
+    cuda::check_cuda_error(::cudaGetDevice(&device_num));
+    return device_num;
+}
+
+/** @brief Set current GPU based on device number.*/
+inline void set_device(int device_num) {
+    cuda::check_cuda_error(::cudaSetDevice(device_num));
+}
+
 }  // namespace caribou

pyopensn/CMakeLists.txt

@@ -34,6 +34,9 @@ target_include_directories(__init__
     ${PROJECT_SOURCE_DIR}/external
     ${MPI4PY_INCLUDE_DIR}
 )
+if (OPENSN_WITH_CUDA)
+    target_compile_definitions(__init__ PRIVATE __OPENSN_USE_CUDA__)
+endif()
 target_link_libraries(__init__
     PRIVATE
     libopensn

pyopensn/py_module.cc

@@ -14,6 +14,11 @@ PYBIND11_MODULE(pyopensn, pyopensn)
   // Metadata
   pyopensn.doc() = "Python interface for OpenSn.";
   pyopensn.attr("__version__") = PROJECT_VERSION;
+#ifdef __OPENSN_USE_CUDA__
+  pyopensn.attr("can_support_gpus") = true;
+#else
+  pyopensn.attr("can_support_gpus") = false;
+#endif // __OPENSN_USE_CUDA__
   // Environment
   if (not PyEnv::p_default_env)
   {
@@ -32,4 +37,7 @@ PYBIND11_MODULE(pyopensn, pyopensn)
   py_ffunc(pyopensn);
   py_response(pyopensn);
   py_solver(pyopensn);
+#ifdef __OPENSN_USE_CUDA__
+  py_device(pyopensn);
+#endif // __OPENSN_USE_CUDA__
 }

python/CMakeLists.txt

@@ -14,6 +14,10 @@ find_package(pybind11 CONFIG REQUIRED)
 
 # libopensnpython
 file(GLOB_RECURSE LIBOPENSN_PY_SRCS CONFIGURE_DEPENDS lib/*.cc)
+if (OPENSN_WITH_CUDA)
+    file(GLOB_RECURSE LIBOPENSN_PY_CUDA_SRCS CONFIGURE_DEPENDS lib/*.cu)
+    list(APPEND LIBOPENSN_PY_SRCS ${LIBOPENSN_PY_CUDA_SRCS})
+endif()
 add_library(libopensnpy STATIC ${LIBOPENSN_PY_SRCS})
 target_include_directories(libopensnpy
     PRIVATE
@@ -32,6 +36,9 @@ target_link_libraries(libopensnpy
     pybind11::embed
 )
 target_compile_options(libopensnpy PRIVATE ${OPENSN_CXX_FLAGS})
+if (OPENSN_WITH_CUDA)
+    target_compile_definitions(libopensnpy PRIVATE __OPENSN_USE_CUDA__)
+endif()
 set_target_properties(
     libopensnpy
     PROPERTIES

python/lib/aquad.cc

@@ -354,14 +354,14 @@ WrapLebedevQuadrature(py::module& aquad)
     "LebedevQuadrature",
     R"(
     Lebedev quadrature for spherical integration.
-    
+
     This quadrature provides high-order accuracy for spherical integration with
     symmetric distribution of points on the sphere.
 
     Wrapper of :cpp:class:`opensn::LebedevQuadrature`.
     )"
   );
-  
+
   lebedev_quadrature.def(
     py::init(
       [](py::kwargs& params)

python/lib/device.cu

@@ -0,0 +1,54 @@
+// SPDX-FileCopyrightText: 2025 The OpenSn Authors <https://open-sn.github.io/opensn/>
+// SPDX-License-Identifier: MIT
+
+#include "python/lib/py_wrappers.h"
+#include "caribou/caribou.h"
+
+namespace crb = caribou;
+
+namespace opensn
+{
+
+// Wrap device settings
+void
+WrapDeviceSettings(py::module& device)
+{
+  // clang-format off
+  device.def(
+    "get_device_count",
+    []() {
+      return crb::get_num_gpus();
+    },
+    "Get the number of GPU devices visible to the current MPI rank."
+  );
+  device.def(
+    "get_current_device",
+    []() {
+      return crb::get_current_device();
+    },
+    "Get the current GPU device number for the current MPI rank."
+  );
+  device.def(
+    "set_device",
+    [](int device_num) {
+      crb::set_device(device_num);
+    },
+    R"(
+    Set GPU device for the current MPI rank using device number
+
+    The device number must range from 0 up to (but not including) ``get_device_count()``.
+    )",
+    py::arg("device_num")
+  );
+  // clang-format on
+}
+
+// Wrap the device setting components of OpenSn
+void
+py_device(py::module& pyopensn)
+{
+  py::module device = pyopensn.def_submodule("device", "GPU settings module.");
+  WrapDeviceSettings(device);
+}
+
+} // namespace opensn

python/lib/py_app.cc

@@ -26,6 +26,11 @@ PyApp::PyApp(const mpi::Communicator& comm)
   py::exec("rank = " + std::to_string(comm.rank()));
   py::exec("size = " + std::to_string(comm.size()));
   py::exec("opensn_console = True");
+#ifdef __OPENSN_USE_CUDA__
+  py::exec("can_support_gpus = True");
+#else
+  py::exec("can_support_gpus = False");
+#endif // __OPENSN_USE_CUDA__
 
   console.BindBarrier(comm);
 
@@ -64,6 +69,10 @@ PyApp::PyApp(const mpi::Communicator& comm)
   Console::BindModule(WrapNLKEigen);
   Console::BindModule(WrapPIteration);
   Console::BindModule(WrapDiscreteOrdinatesKEigenAcceleration);
+
+#ifdef __OPENSN_USE_CUDA__
+  Console::BindModule(WrapDeviceSettings);
+#endif // __OPENSN_USE_CUDA__
 }
 
 int

python/lib/py_wrappers.h

@@ -238,4 +238,10 @@ void WrapVolumetricSource(py::module& src);
 void py_xs(py::module& pyopensn);
 void WrapMultiGroupXS(py::module& xs);
 
-} // namespace opensn
+#ifdef __OPENSN_USE_CUDA__
+/// Wrap the CUDA/ROCm components of OpenSn.
+void py_device(py::module& pyopensn);
+void WrapDeviceSettings(py::module& device);
+#endif // __OPENSN_USE_CUDA__
+
+} // namespace opensn