sundials4py: add cuda support

bindings/sundials4py/CMakeLists.txt

@@ -116,6 +116,11 @@ set(sundials_SOURCES
     sunnonlinsol/sunnonlinsol_newton.cpp
     test/sundials4py_test.cpp)
 
+if(SUNDIALS_ENABLE_NVECTOR_CUDA)
+  list(APPEND sundials_SOURCES nvector/nvector_cuda.cpp
+       sunmemory/sunmemory_cuda.cpp)
+endif()
+
 # Create the Python sundials library
 nanobind_add_module(sundials4py STABLE_ABI NB_STATIC ${sundials_SOURCES})
 
@@ -150,5 +155,10 @@ target_link_libraries(
           sundials_sundomeigestpower
           sundials_core)
 
+if(SUNDIALS_ENABLE_NVECTOR_CUDA)
+  target_link_libraries(sundials4py PRIVATE sundials_nveccuda CUDA::cudart)
+  target_include_directories(sundials4py PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+endif()
+
 # Install directive for scikit-build-core
 install(TARGETS sundials4py LIBRARY DESTINATION .)

bindings/sundials4py/nvector/generate.yaml

@@ -104,3 +104,23 @@ modules:
     # Getting and setting array pointers requires us to do something custom
     - "^N_VGetSubvectorArrayPointer_ManyVector$"
     - "^N_VSetSubvectorArrayPointer_ManyVector$"
+  nvector_cuda:
+    path: nvector/nvector_cuda_generated.hpp
+    headers:
+    - ../../include/nvector/nvector_cuda.h
+    fn_exclude_by_name__regex:
+    # Host/device pointer access and constructors need custom ownership and
+    # pointer-address handling in Python.
+    - "^N_VNewEmpty_Cuda$"
+    - "^N_VNew_Cuda$"
+    - "^N_VNewManaged_Cuda$"
+    - "^N_VNewWithMemHelp_Cuda$"
+    - "^N_VMake_Cuda$"
+    - "^N_VMakeManaged_Cuda$"
+    - "^N_VSetHostArrayPointer_Cuda$"
+    - "^N_VSetDeviceArrayPointer_Cuda$"
+    - "^N_VGetHostArrayPointer_Cuda$"
+    - "^N_VGetDeviceArrayPointer_Cuda$"
+    # Kernel execution policies are C++ CUDA objects that are not currently
+    # exposed through sundials4py.
+    - "^N_VSetKernelExecPolicy_Cuda$"

bindings/sundials4py/nvector/nvector_cuda.cpp

@@ -0,0 +1,191 @@
+#include "sundials4py.hpp"
+
+#include <cstdint>
+
+#include <nvector/nvector_cuda.h>
+#include <sundials/sundials_nvector.hpp>
+
+#include "sundials/sundials_classview.hpp"
+
+namespace nb = nanobind;
+using namespace sundials::experimental;
+
+namespace sundials4py {
+
+namespace {
+
+using CudaArray1d =
+  nb::ndarray<sunrealtype, nb::device::cuda, nb::ndim<1>, nb::c_contig>;
+
+void check_length(sunindextype vec_length, sundials4py::Array1d data)
+{
+  if (data.shape(0) != static_cast<size_t>(vec_length))
+  {
+    throw sundials4py::error_returned(
+      "Array shape does not match vector length");
+  }
+}
+
+void check_length(sunindextype vec_length, CudaArray1d data)
+{
+  if (data.shape(0) != static_cast<size_t>(vec_length))
+  {
+    throw sundials4py::error_returned(
+      "CUDA array shape does not match vector length");
+  }
+}
+
+std::shared_ptr<std::remove_pointer_t<N_Vector>> wrap_nvector(N_Vector v)
+{
+  return our_make_shared<std::remove_pointer_t<N_Vector>, N_VectorDeleter>(v);
+}
+
+} // namespace
+
+void bind_nvector_cuda(nb::module_& m)
+{
+#include "nvector_cuda_generated.hpp"
+
+  m.def(
+    "N_VNewEmpty_Cuda",
+    [](SUNContext sunctx) -> std::shared_ptr<std::remove_pointer_t<N_Vector>>
+    { return wrap_nvector(N_VNewEmpty_Cuda(sunctx)); },
+    nb::arg("sunctx"), nb::keep_alive<0, 1>());
+
+  m.def(
+    "N_VNew_Cuda",
+    [](sunindextype vec_length,
+       SUNContext sunctx) -> std::shared_ptr<std::remove_pointer_t<N_Vector>>
+    { return wrap_nvector(N_VNew_Cuda(vec_length, sunctx)); },
+    nb::arg("vec_length"), nb::arg("sunctx"), nb::keep_alive<0, 2>());
+
+  m.def(
+    "N_VNewManaged_Cuda",
+    [](sunindextype vec_length,
+       SUNContext sunctx) -> std::shared_ptr<std::remove_pointer_t<N_Vector>>
+    { return wrap_nvector(N_VNewManaged_Cuda(vec_length, sunctx)); },
+    nb::arg("vec_length"), nb::arg("sunctx"), nb::keep_alive<0, 2>());
+
+  m.def(
+    "N_VNewWithMemHelp_Cuda",
+    [](sunindextype vec_length, sunbooleantype use_managed_mem,
+       SUNMemoryHelper helper,
+       SUNContext sunctx) -> std::shared_ptr<std::remove_pointer_t<N_Vector>>
+    {
+      return wrap_nvector(
+        N_VNewWithMemHelp_Cuda(vec_length, use_managed_mem, helper, sunctx));
+    },
+    nb::arg("vec_length"), nb::arg("use_managed_mem"), nb::arg("helper"),
+    nb::arg("sunctx"), nb::keep_alive<0, 3>(), nb::keep_alive<0, 4>());
+
+  m.def(
+    "N_VMake_Cuda",
+    [](sunindextype vec_length, sundials4py::Array1d h_vdata_1d,
+       std::uintptr_t d_vdata,
+       SUNContext sunctx) -> std::shared_ptr<std::remove_pointer_t<N_Vector>>
+    {
+      check_length(vec_length, h_vdata_1d);
+      return wrap_nvector(N_VMake_Cuda(vec_length, h_vdata_1d.data(),
+                                       reinterpret_cast<sunrealtype*>(d_vdata),
+                                       sunctx));
+    },
+    nb::arg("vec_length"), nb::arg("h_vdata_1d"), nb::arg("d_vdata"),
+    nb::arg("sunctx"), nb::keep_alive<0, 2>(), nb::keep_alive<0, 4>());
+
+  m.def(
+    "N_VMake_Cuda",
+    [](sunindextype vec_length, sundials4py::Array1d h_vdata_1d,
+       CudaArray1d d_vdata_1d,
+       SUNContext sunctx) -> std::shared_ptr<std::remove_pointer_t<N_Vector>>
+    {
+      check_length(vec_length, h_vdata_1d);
+      check_length(vec_length, d_vdata_1d);
+      return wrap_nvector(
+        N_VMake_Cuda(vec_length, h_vdata_1d.data(), d_vdata_1d.data(), sunctx));
+    },
+    nb::arg("vec_length"), nb::arg("h_vdata_1d").noconvert(),
+    nb::arg("d_vdata_1d").noconvert(), nb::arg("sunctx"),
+    nb::keep_alive<0, 2>(), nb::keep_alive<0, 3>(), nb::keep_alive<0, 4>());
+
+  m.def(
+    "N_VMakeManaged_Cuda",
+    [](sunindextype vec_length, std::uintptr_t vdata,
+       SUNContext sunctx) -> std::shared_ptr<std::remove_pointer_t<N_Vector>>
+    {
+      return wrap_nvector(
+        N_VMakeManaged_Cuda(vec_length, reinterpret_cast<sunrealtype*>(vdata),
+                            sunctx));
+    },
+    nb::arg("vec_length"), nb::arg("vdata"), nb::arg("sunctx"),
+    nb::keep_alive<0, 3>());
+
+  m.def(
+    "N_VSetHostArrayPointer_Cuda",
+    [](sundials4py::Array1d h_vdata_1d, N_Vector v)
+    {
+      check_length(N_VGetLength_Cuda(v), h_vdata_1d);
+      N_VSetHostArrayPointer_Cuda(h_vdata_1d.data(), v);
+    },
+    nb::arg("h_vdata_1d"), nb::arg("v"));
+
+  m.def(
+    "N_VSetDeviceArrayPointer_Cuda",
+    [](std::uintptr_t d_vdata, N_Vector v) {
+      N_VSetDeviceArrayPointer_Cuda(reinterpret_cast<sunrealtype*>(d_vdata), v);
+    },
+    nb::arg("d_vdata"), nb::arg("v"));
+
+  m.def(
+    "N_VSetDeviceArrayPointer_Cuda",
+    [](CudaArray1d d_vdata_1d, N_Vector v)
+    {
+      check_length(N_VGetLength_Cuda(v), d_vdata_1d);
+      N_VSetDeviceArrayPointer_Cuda(d_vdata_1d.data(), v);
+    },
+    nb::arg("d_vdata_1d").noconvert(), nb::arg("v"), nb::keep_alive<2, 1>());
+
+  m.def(
+    "N_VGetHostArrayPointer_Cuda",
+    [](N_Vector x)
+    {
+      auto ptr = N_VGetHostArrayPointer_Cuda(x);
+      if (!ptr)
+      {
+        throw sundials4py::error_returned("Failed to get host array pointer");
+      }
+      auto owner = nb::find(x);
+      size_t shape[1]{static_cast<size_t>(N_VGetLength_Cuda(x))};
+      return sundials4py::Array1d(ptr, 1, shape, owner);
+    },
+    nb::rv_policy::reference);
+
+  m.def(
+    "N_VGetDeviceArrayPointer_Cuda",
+    [](N_Vector x) -> std::uintptr_t
+    {
+      auto ptr = N_VGetDeviceArrayPointer_Cuda(x);
+      if (!ptr)
+      {
+        throw sundials4py::error_returned("Failed to get device array pointer");
+      }
+      return reinterpret_cast<std::uintptr_t>(ptr);
+    },
+    nb::arg("x"));
+
+  m.def(
+    "N_VGetDeviceArray_Cuda",
+    [](N_Vector x)
+    {
+      auto ptr = N_VGetDeviceArrayPointer_Cuda(x);
+      if (!ptr)
+      {
+        throw sundials4py::error_returned("Failed to get device array pointer");
+      }
+      auto owner = nb::find(x);
+      size_t shape[1]{static_cast<size_t>(N_VGetLength_Cuda(x))};
+      return CudaArray1d(ptr, 1, shape, owner);
+    },
+    nb::rv_policy::reference);
+}
+
+} // namespace sundials4py

bindings/sundials4py/nvector/nvector_cuda_generated.hpp

@@ -0,0 +1,55 @@
+// #ifndef _NVECTOR_CUDA_H
+//
+// #ifdef __cplusplus
+// #endif
+//
+
+auto pyClass_N_VectorContent_Cuda =
+  nb::class_<_N_VectorContent_Cuda>(m, "_N_VectorContent_Cuda", "")
+    .def(nb::init<>()) // implicit default constructor
+  ;
+
+m.def("N_VIsManagedMemory_Cuda", N_VIsManagedMemory_Cuda, nb::arg("x"));
+
+m.def("N_VCopyToDevice_Cuda", N_VCopyToDevice_Cuda, nb::arg("v"));
+
+m.def("N_VCopyFromDevice_Cuda", N_VCopyFromDevice_Cuda, nb::arg("v"));
+
+m.def("N_VEnableFusedOps_Cuda", N_VEnableFusedOps_Cuda, nb::arg("v"),
+      nb::arg("tf"));
+
+m.def("N_VEnableLinearCombination_Cuda", N_VEnableLinearCombination_Cuda,
+      nb::arg("v"), nb::arg("tf"));
+
+m.def("N_VEnableScaleAddMulti_Cuda", N_VEnableScaleAddMulti_Cuda, nb::arg("v"),
+      nb::arg("tf"));
+
+m.def("N_VEnableDotProdMulti_Cuda", N_VEnableDotProdMulti_Cuda, nb::arg("v"),
+      nb::arg("tf"));
+
+m.def("N_VEnableLinearSumVectorArray_Cuda", N_VEnableLinearSumVectorArray_Cuda,
+      nb::arg("v"), nb::arg("tf"));
+
+m.def("N_VEnableScaleVectorArray_Cuda", N_VEnableScaleVectorArray_Cuda,
+      nb::arg("v"), nb::arg("tf"));
+
+m.def("N_VEnableConstVectorArray_Cuda", N_VEnableConstVectorArray_Cuda,
+      nb::arg("v"), nb::arg("tf"));
+
+m.def("N_VEnableWrmsNormVectorArray_Cuda", N_VEnableWrmsNormVectorArray_Cuda,
+      nb::arg("v"), nb::arg("tf"));
+
+m.def("N_VEnableWrmsNormMaskVectorArray_Cuda",
+      N_VEnableWrmsNormMaskVectorArray_Cuda, nb::arg("v"), nb::arg("tf"));
+
+m.def("N_VEnableScaleAddMultiVectorArray_Cuda",
+      N_VEnableScaleAddMultiVectorArray_Cuda, nb::arg("v"), nb::arg("tf"));
+
+m.def("N_VEnableLinearCombinationVectorArray_Cuda",
+      N_VEnableLinearCombinationVectorArray_Cuda, nb::arg("v"), nb::arg("tf"));
+// #ifdef __cplusplus
+//
+// #endif
+//
+// #endif
+//

bindings/sundials4py/sundials4py.cpp

@@ -40,8 +40,14 @@ void bind_kinsol(nb::module_& m);
 
 void bind_nvector_serial(nb::module_& m);
 void bind_nvector_manyvector(nb::module_& m);
+#ifdef SUNDIALS_NVECTOR_CUDA
+void bind_nvector_cuda(nb::module_& m);
+#endif
 
 void bind_sumemoryhelper_sys(nb::module_& m);
+#ifdef SUNDIALS_NVECTOR_CUDA
+void bind_sunmemoryhelper_cuda(nb::module_& m);
+#endif
 
 void bind_sunadaptcontroller_imexgus(nb::module_& m);
 void bind_sunadaptcontroller_mrihtol(nb::module_& m);
@@ -119,6 +125,9 @@ NB_MODULE(sundials4py, m)
 
   sundials4py::bind_nvector_serial(core_m);
   sundials4py::bind_nvector_manyvector(core_m);
+#ifdef SUNDIALS_NVECTOR_CUDA
+  sundials4py::bind_nvector_cuda(core_m);
+#endif
 
   sundials4py::bind_sunadaptcontroller_imexgus(core_m);
   sundials4py::bind_sunadaptcontroller_mrihtol(core_m);
@@ -141,8 +150,11 @@ NB_MODULE(sundials4py, m)
   sundials4py::bind_sunmatrix_sparse(core_m);
 
   sundials4py::bind_sumemoryhelper_sys(core_m);
+#ifdef SUNDIALS_NVECTOR_CUDA
+  sundials4py::bind_sunmemoryhelper_cuda(core_m);
+#endif
 
   sundials4py::bind_sunnonlinsol_fixedpoint(core_m);
   sundials4py::bind_sunnonlinsol_newton(core_m);
   sundials4py::bind_sunnonlinsol_auto(core_m);
-}
+}

bindings/sundials4py/sunmemory/generate.yaml

@@ -38,3 +38,7 @@ modules:
     path: sunmemory/sunmemory_system_generated.hpp
     headers:
     - ../../include/sunmemory/sunmemory_system.h
+  sunmemory_cuda:
+    path: sunmemory/sunmemory_cuda_generated.hpp
+    headers:
+    - ../../include/sunmemory/sunmemory_cuda.h

bindings/sundials4py/sunmemory/sunmemory_cuda.cpp

@@ -0,0 +1,33 @@
+/* -----------------------------------------------------------------
+ * Programmer(s): Cody J. Balos @ LLNL
+ * -----------------------------------------------------------------
+ * SUNDIALS Copyright Start
+ * Copyright (c) 2025-2026, Lawrence Livermore National Security,
+ * University of Maryland Baltimore County, and the SUNDIALS contributors.
+ * Copyright (c) 2013-2025, Lawrence Livermore National Security
+ * and Southern Methodist University.
+ * Copyright (c) 2002-2013, Lawrence Livermore National Security.
+ * All rights reserved.
+ *
+ * See the top-level LICENSE and NOTICE files for details.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ * SUNDIALS Copyright End
+ * -----------------------------------------------------------------*/
+
+#include "sundials4py.hpp"
+
+#include <sundials/sundials_core.hpp>
+#include <sunmemory/sunmemory_cuda.h>
+
+namespace nb = nanobind;
+using namespace sundials::experimental;
+
+namespace sundials4py {
+
+void bind_sunmemoryhelper_cuda(nb::module_& m)
+{
+#include "sunmemory_cuda_generated.hpp"
+}
+
+} // namespace sundials4py