From 8330478457e65c0633a27956032c5925f6b458dc Mon Sep 17 00:00:00 2001 From: Cody Balos Date: Tue, 30 Jun 2026 16:43:59 -0700 Subject: [PATCH 1/3] add cuda support --- bindings/sundials4py/CMakeLists.txt | 10 + bindings/sundials4py/nvector/generate.yaml | 20 ++ bindings/sundials4py/nvector/nvector_cuda.cpp | 191 +++++++++++++++++ .../nvector/nvector_cuda_generated.hpp | 57 +++++ bindings/sundials4py/sundials4py.cpp | 14 +- bindings/sundials4py/sunmemory/generate.yaml | 4 + .../sundials4py/sunmemory/sunmemory_cuda.cpp | 33 +++ .../sunmemory/sunmemory_cuda_generated.hpp | 21 ++ bindings/sundials4py/test/test_nvector.py | 89 ++++++++ .../sundials4py/test/test_sunmemoryhelper.py | 10 + examples/python/arkode/ark_heat1D_cuda.py | 197 ++++++++++++++++++ 11 files changed, 645 insertions(+), 1 deletion(-) create mode 100644 bindings/sundials4py/nvector/nvector_cuda.cpp create mode 100644 bindings/sundials4py/nvector/nvector_cuda_generated.hpp create mode 100644 bindings/sundials4py/sunmemory/sunmemory_cuda.cpp create mode 100644 bindings/sundials4py/sunmemory/sunmemory_cuda_generated.hpp create mode 100644 examples/python/arkode/ark_heat1D_cuda.py diff --git a/bindings/sundials4py/CMakeLists.txt b/bindings/sundials4py/CMakeLists.txt index 7b607f2073..4982955914 100644 --- a/bindings/sundials4py/CMakeLists.txt +++ b/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 .) diff --git a/bindings/sundials4py/nvector/generate.yaml b/bindings/sundials4py/nvector/generate.yaml index bef1cd0e0a..f5e5a768c0 100644 --- a/bindings/sundials4py/nvector/generate.yaml +++ b/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$" diff --git a/bindings/sundials4py/nvector/nvector_cuda.cpp b/bindings/sundials4py/nvector/nvector_cuda.cpp new file mode 100644 index 0000000000..955447fa7e --- /dev/null +++ b/bindings/sundials4py/nvector/nvector_cuda.cpp @@ -0,0 +1,191 @@ +#include "sundials4py.hpp" + +#include + +#include +#include + +#include "sundials/sundials_classview.hpp" + +namespace nb = nanobind; +using namespace sundials::experimental; + +namespace sundials4py { + +namespace { + +using CudaArray1d = + nb::ndarray, nb::c_contig>; + +void check_length(sunindextype vec_length, sundials4py::Array1d data) +{ + if (data.shape(0) != static_cast(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(vec_length)) + { + throw sundials4py::error_returned( + "CUDA array shape does not match vector length"); + } +} + +std::shared_ptr> wrap_nvector(N_Vector v) +{ + return our_make_shared, 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> + { 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> + { 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> + { 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> + { + 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> + { + check_length(vec_length, h_vdata_1d); + return wrap_nvector(N_VMake_Cuda(vec_length, h_vdata_1d.data(), + reinterpret_cast(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> + { + 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> + { + return wrap_nvector( + N_VMakeManaged_Cuda(vec_length, reinterpret_cast(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(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(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(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(N_VGetLength_Cuda(x))}; + return CudaArray1d(ptr, 1, shape, owner); + }, + nb::rv_policy::reference); +} + +} // namespace sundials4py diff --git a/bindings/sundials4py/nvector/nvector_cuda_generated.hpp b/bindings/sundials4py/nvector/nvector_cuda_generated.hpp new file mode 100644 index 0000000000..bdfde029f8 --- /dev/null +++ b/bindings/sundials4py/nvector/nvector_cuda_generated.hpp @@ -0,0 +1,57 @@ +// #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_VGetLength_Cuda", N_VGetLength_Cuda, nb::arg("x")); + +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 +// diff --git a/bindings/sundials4py/sundials4py.cpp b/bindings/sundials4py/sundials4py.cpp index d23b9bd2b1..392dc322e3 100644 --- a/bindings/sundials4py/sundials4py.cpp +++ b/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); -} \ No newline at end of file +} diff --git a/bindings/sundials4py/sunmemory/generate.yaml b/bindings/sundials4py/sunmemory/generate.yaml index 5cc2e01d8d..a80d2c9452 100644 --- a/bindings/sundials4py/sunmemory/generate.yaml +++ b/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 diff --git a/bindings/sundials4py/sunmemory/sunmemory_cuda.cpp b/bindings/sundials4py/sunmemory/sunmemory_cuda.cpp new file mode 100644 index 0000000000..29e599370f --- /dev/null +++ b/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 +#include + +namespace nb = nanobind; +using namespace sundials::experimental; + +namespace sundials4py { + +void bind_sunmemoryhelper_cuda(nb::module_& m) +{ +#include "sunmemory_cuda_generated.hpp" +} + +} // namespace sundials4py diff --git a/bindings/sundials4py/sunmemory/sunmemory_cuda_generated.hpp b/bindings/sundials4py/sunmemory/sunmemory_cuda_generated.hpp new file mode 100644 index 0000000000..43c2a9a4c5 --- /dev/null +++ b/bindings/sundials4py/sunmemory/sunmemory_cuda_generated.hpp @@ -0,0 +1,21 @@ +// #ifndef _SUNDIALS_CUDAMEMORY_H +// +// #ifdef __cplusplus +// #endif +// + +m.def( + "SUNMemoryHelper_Cuda", + [](SUNContext sunctx) -> std::shared_ptr> + { + auto helper = SUNMemoryHelper_Cuda(sunctx); + return our_make_shared, + SUNMemoryHelperDeleter>(helper); + }, + nb::arg("sunctx"), nb::keep_alive<0, 1>()); +// #ifdef __cplusplus +// +// #endif +// +// #endif +// diff --git a/bindings/sundials4py/test/test_nvector.py b/bindings/sundials4py/test/test_nvector.py index 1afcfa1428..d77c64e282 100644 --- a/bindings/sundials4py/test/test_nvector.py +++ b/bindings/sundials4py/test/test_nvector.py @@ -78,6 +78,95 @@ def test_make_nvector(vector_type, sunctx): assert_allclose(N_VGetArrayPointer(nvec), [5.0, 4.0, 3.0, 2.0, 1.0]) +@pytest.mark.skipif("N_VNew_Cuda" not in globals(), reason="CUDA bindings are not enabled") +def test_create_nvector_cuda(sunctx): + try: + nvec = N_VNew_Cuda(5, sunctx) + except RuntimeError as err: + pytest.skip(f"CUDA vector allocation failed: {err}") + + if nvec is None: + pytest.skip("CUDA vector allocation failed") + + assert N_VGetLength(nvec) == 5 + assert N_VGetLength_Cuda(nvec) == 5 + assert N_VGetDeviceArrayPointer_Cuda(nvec) != 0 + + arr = N_VGetHostArrayPointer_Cuda(nvec) + arr[:] = np.array([5.0, 4.0, 3.0, 2.0, 1.0], dtype=sunrealtype) + + N_VCopyToDevice_Cuda(nvec) + N_VConst(2.0, nvec) + N_VCopyFromDevice_Cuda(nvec) + + assert_allclose(arr, 2.0) + + +def _torch_dtype(): + import torch + + if sunrealtype == np.float32: + return torch.float32 + if sunrealtype == np.float64: + return torch.float64 + return torch.longdouble + + +@pytest.mark.skipif("N_VMake_Cuda" not in globals(), reason="CUDA bindings are not enabled") +def test_make_nvector_cuda_cupy_array(sunctx): + cupy = pytest.importorskip("cupy") + + h_arr = np.zeros(5, dtype=sunrealtype) + d_arr = cupy.arange(5, dtype=sunrealtype) + + nvec = N_VMake_Cuda(5, h_arr, d_arr, sunctx) + N_VConst(3.0, nvec) + N_VCopyFromDevice_Cuda(nvec) + + assert_allclose(cupy.asnumpy(d_arr), 3.0) + assert_allclose(h_arr, 3.0) + + view = cupy.from_dlpack(N_VGetDeviceArray_Cuda(nvec)) + assert_allclose(cupy.asnumpy(view), 3.0) + + +@pytest.mark.skipif("N_VMake_Cuda" not in globals(), reason="CUDA bindings are not enabled") +def test_make_nvector_cuda_torch_tensor(sunctx): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA is not available") + + h_arr = np.zeros(5, dtype=sunrealtype) + d_arr = torch.arange(5, device="cuda", dtype=_torch_dtype()) + + nvec = N_VMake_Cuda(5, h_arr, d_arr, sunctx) + N_VConst(4.0, nvec) + N_VCopyFromDevice_Cuda(nvec) + + assert_allclose(d_arr.cpu().numpy(), 4.0) + assert_allclose(h_arr, 4.0) + + view = torch.utils.dlpack.from_dlpack(N_VGetDeviceArray_Cuda(nvec)) + assert_allclose(view.cpu().numpy(), 4.0) + + +@pytest.mark.skipif( + "N_VSetDeviceArrayPointer_Cuda" not in globals(), reason="CUDA bindings are not enabled" +) +def test_set_nvector_cuda_torch_tensor(sunctx): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA is not available") + + nvec = N_VNew_Cuda(5, sunctx) + d_arr = torch.arange(5, device="cuda", dtype=_torch_dtype()) + + N_VSetDeviceArrayPointer_Cuda(d_arr, nvec) + N_VConst(5.0, nvec) + + assert_allclose(d_arr.cpu().numpy(), 5.0) + + # Test an operation that involves vector arrays @pytest.mark.parametrize("vector_type", ["serial"]) def test_nvlinearcombination(vector_type, sunctx): diff --git a/bindings/sundials4py/test/test_sunmemoryhelper.py b/bindings/sundials4py/test/test_sunmemoryhelper.py index 88728e079b..e259482ae7 100644 --- a/bindings/sundials4py/test/test_sunmemoryhelper.py +++ b/bindings/sundials4py/test/test_sunmemoryhelper.py @@ -25,3 +25,13 @@ def test_create_memory_helper_sys(sunctx): mem_helper = SUNMemoryHelper_Sys(sunctx) # noqa: F405 assert mem_helper is not None + + +@pytest.mark.skipif( + "SUNMemoryHelper_Cuda" not in globals(), reason="CUDA bindings are not enabled" +) +def test_create_memory_helper_cuda(sunctx): + mem_helper = SUNMemoryHelper_Cuda(sunctx) # noqa: F405 + if mem_helper is None: + pytest.skip("CUDA memory helper creation failed") + assert mem_helper is not None diff --git a/examples/python/arkode/ark_heat1D_cuda.py b/examples/python/arkode/ark_heat1D_cuda.py new file mode 100644 index 0000000000..41c825d5f7 --- /dev/null +++ b/examples/python/arkode/ark_heat1D_cuda.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------- +# Programmer(s): Cody J. Balos +# ----------------------------------------------------------------- +# 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 +# ----------------------------------------------------------------- +# This example is a copy of examples/arkode/C_serial/ark_heat1D.c, +# but ported to Python to use sundials4py as well as CUDA through +# either CuPy or PyTorch. +# +# The following test simulates a simple 1D heat equation, +# u_t = k*u_xx + f +# for t in [0, 10], x in [0, 1], with initial conditions +# u(0,x) = 0 +# Dirichlet boundary conditions, i.e. +# u_t(t,0) = u_t(t,1) = 0, +# and a point-source heating term, +# f = 0.01 for x=0.5. +# +# The spatial derivatives are computed using second-order +# centered differences, with the data distributed over N points +# on a uniform spatial grid. +# +# This program solves the problem with either an ERK or DIRK +# method. For the DIRK method, we use a Newton iteration with +# the SUNLinSol_PCG linear solver, and a user-supplied Jacobian-vector +# product routine. +# +# 100 outputs are printed at equal intervals, and run statistics +# are printed at the end. +# ----------------------------------------------------------------- + +import numpy as np +import sundials4py.arkode as ark +import sundials4py.core as sun + + +class Heat1DCudaProblem: + def __init__(self, backend, n=101, k=0.01): + self.backend = backend + self.n = n + self.k = k + self.dx = 1.0 / (n - 1) + self.isource = n // 2 + + def device_array(self, nvec): + return self.backend.from_dlpack(sun.N_VGetDeviceArray_Cuda(nvec)) + + def set_init_cond(self, yvec): + y = self.device_array(yvec) + y[:] = 0.0 + + def f(self, t, yvec, ydotvec, user_data): + y = self.device_array(yvec) + ydot = self.device_array(ydotvec) + + ydot[:] = 0.0 + c1 = self.k / self.dx / self.dx + c2 = -2.0 * self.k / self.dx / self.dx + ydot[1:-1] = c1 * y[:-2] + c2 * y[1:-1] + c1 * y[2:] + ydot[0] = 0.0 + ydot[-1] = 0.0 + ydot[self.isource] += 0.01 / self.dx + self.backend.synchronize() + return 0 + + +class CupyBackend: + name = "cupy" + + def __init__(self): + import cupy + + self.xp = cupy + + def zeros(self, n): + return self.xp.zeros(n, dtype=sun.sunrealtype) + + def from_dlpack(self, obj): + return self.xp.from_dlpack(obj) + + def to_numpy(self, array): + return self.xp.asnumpy(array) + + def synchronize(self): + self.xp.cuda.Device().synchronize() + + +class TorchBackend: + name = "torch" + + def __init__(self): + import torch + + if not torch.cuda.is_available(): + raise RuntimeError("PyTorch CUDA is not available") + self.torch = torch + self.dtype = torch.float32 if sun.sunrealtype == np.float32 else torch.float64 + + def zeros(self, n): + return self.torch.zeros(n, device="cuda", dtype=self.dtype) + + def from_dlpack(self, obj): + return self.torch.utils.dlpack.from_dlpack(obj) + + def to_numpy(self, array): + return array.cpu().numpy() + + def synchronize(self): + self.torch.cuda.synchronize() + + +def solve_heat1d(backend): + n = 101 + tf = 1.0 + nt = 10 + reltol = 1e-6 + abstol = 1e-10 + + status, sunctx = sun.SUNContext_Create(sun.SUN_COMM_NULL) + assert status == sun.SUN_SUCCESS + + host_data = np.zeros(n, dtype=sun.sunrealtype) + device_data = backend.zeros(n) + y = sun.N_VMake_Cuda(n, host_data, device_data, sunctx) + + problem = Heat1DCudaProblem(backend, n=n) + problem.set_init_cond(y) + + stepper = ark.ARKStepCreate(problem.f, None, 0.0, y, sunctx) + assert stepper is not None + + status = ark.ARKodeSStolerances(stepper.get(), reltol, abstol) + assert status == ark.ARK_SUCCESS + + status = ark.ARKodeSetMaxNumSteps(stepper.get(), 100000) + assert status == ark.ARK_SUCCESS + + t = 0.0 + tout = tf / nt + + print(f"\n{backend.name} backend") + print(" t ||u||_rms") + print(" -------------------------") + print(f" {t:10.6f} {0.0:10.6f}") + + for _ in range(nt): + status, t = ark.ARKodeEvolve(stepper.get(), tout, y, ark.ARK_NORMAL) + if status != ark.ARK_SUCCESS: + raise RuntimeError(f"ARKodeEvolve failed with status {status}") + + sun.N_VCopyFromDevice_Cuda(y) + rms = np.sqrt(np.dot(host_data, host_data) / n) + print(f" {t:10.6f} {rms:10.6f}") + tout = min(tout + tf / nt, tf) + + return host_data.copy(), backend.to_numpy(device_data) + + +def main(): + if not hasattr(sun, "N_VMake_Cuda"): + raise RuntimeError("sundials4py was not built with CUDA support") + + results = {} + for backend_type in (CupyBackend, TorchBackend): + try: + backend = backend_type() + except ImportError: + continue + except RuntimeError as err: + print(f"Skipping {backend_type.name}: {err}") + continue + + host_data, device_data = solve_heat1d(backend) + np.testing.assert_allclose(host_data, device_data) + results[backend.name] = host_data + + if not results: + raise RuntimeError("Install CuPy or PyTorch with CUDA support to run this example") + + if {"cupy", "torch"} <= results.keys(): + np.testing.assert_allclose(results["cupy"], results["torch"], rtol=1e-6, atol=1e-10) + + +if __name__ == "__main__": + main() From bb28cff927b343eb10f1439e62b4d444595f5242 Mon Sep 17 00:00:00 2001 From: Cody Balos Date: Tue, 30 Jun 2026 17:01:25 -0700 Subject: [PATCH 2/3] add semi-discrete solution for comparison --- examples/python/arkode/ark_heat1D.py | 34 +++++++++++++++++++++ examples/python/arkode/ark_heat1D_cuda.py | 37 ++++++++++++++++++++++- 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/examples/python/arkode/ark_heat1D.py b/examples/python/arkode/ark_heat1D.py index e24963a6d8..d1481119db 100644 --- a/examples/python/arkode/ark_heat1D.py +++ b/examples/python/arkode/ark_heat1D.py @@ -31,6 +31,19 @@ # centered differences, with the data distributed over N points # on a uniform spatial grid. # +# The final solution is checked against the exact solution of this +# semi-discrete ODE system. With zero Dirichlet boundaries, the +# interior finite-difference Laplacian is diagonalized by the discrete +# sine basis. For zero initial data and a constant point source b, the +# interior solution is +# +# u(t) = Phi * diag((exp(lambda_m*t) - 1) / lambda_m) * Phi^T * b, +# +# where Phi contains the orthonormal sine modes and lambda_m are the +# corresponding finite-difference Laplacian eigenvalues. This validates +# the time integration error without introducing error from a continuous +# PDE approximation. +# # This program solves the problem with either an ERK or DIRK # method. For the DIRK method, we use a Newton iteration with # the SUNLinSol_PCG linear solver, and a user-supplied Jacobian-vector @@ -45,6 +58,22 @@ from sundials4py.arkode import * +def exact_semidiscrete_solution(N, k, t): + dx = 1.0 / (N - 1) + i = np.arange(1, N - 1) + m = np.arange(1, N - 1) + + phi = np.sqrt(2.0 / (N - 1)) * np.sin(np.outer(i, m) * np.pi / (N - 1)) + lambdas = -4.0 * k / dx**2 * np.sin(0.5 * m * np.pi / (N - 1)) ** 2 + + source = np.zeros(N - 2) + source[(N // 2) - 1] = 0.01 / dx + + u = np.zeros(N) + u[1:-1] = phi @ (((np.exp(lambdas * t) - 1.0) / lambdas) * (phi.T @ source)) + return u + + class Heat1DProblem: def __init__(self, N, k): self.N = N @@ -165,6 +194,11 @@ def main(): break print(" -------------------------") + uexact = exact_semidiscrete_solution(N, k, Tf) + max_error = np.max(np.abs(yarr - uexact)) + print(f"\nFinal max error vs exact semi-discrete solution = {max_error:.6e}") + np.testing.assert_allclose(yarr, uexact, rtol=1e-4, atol=1e-8) + # Print statistics status, nst = ARKodeGetNumSteps(ark.get()) assert status == ARK_SUCCESS diff --git a/examples/python/arkode/ark_heat1D_cuda.py b/examples/python/arkode/ark_heat1D_cuda.py index 41c825d5f7..6f2686a626 100644 --- a/examples/python/arkode/ark_heat1D_cuda.py +++ b/examples/python/arkode/ark_heat1D_cuda.py @@ -32,6 +32,19 @@ # centered differences, with the data distributed over N points # on a uniform spatial grid. # +# The final solution is checked against the exact solution of this +# semi-discrete ODE system. With zero Dirichlet boundaries, the +# interior finite-difference Laplacian is diagonalized by the discrete +# sine basis. For zero initial data and a constant point source b, the +# interior solution is +# +# u(t) = Phi * diag((exp(lambda_m*t) - 1) / lambda_m) * Phi^T * b, +# +# where Phi contains the orthonormal sine modes and lambda_m are the +# corresponding finite-difference Laplacian eigenvalues. This validates +# the time integration error without introducing error from a continuous +# PDE approximation. +# # This program solves the problem with either an ERK or DIRK # method. For the DIRK method, we use a Newton iteration with # the SUNLinSol_PCG linear solver, and a user-supplied Jacobian-vector @@ -46,6 +59,22 @@ import sundials4py.core as sun +def exact_semidiscrete_solution(n, k, t): + dx = 1.0 / (n - 1) + i = np.arange(1, n - 1) + m = np.arange(1, n - 1) + + phi = np.sqrt(2.0 / (n - 1)) * np.sin(np.outer(i, m) * np.pi / (n - 1)) + lambdas = -4.0 * k / dx**2 * np.sin(0.5 * m * np.pi / (n - 1)) ** 2 + + source = np.zeros(n - 2) + source[(n // 2) - 1] = 0.01 / dx + + u = np.zeros(n) + u[1:-1] = phi @ (((np.exp(lambdas * t) - 1.0) / lambdas) * (phi.T @ source)) + return u + + class Heat1DCudaProblem: def __init__(self, backend, n=101, k=0.01): self.backend = backend @@ -123,6 +152,7 @@ def synchronize(self): def solve_heat1d(backend): n = 101 + k = 0.01 tf = 1.0 nt = 10 reltol = 1e-6 @@ -135,7 +165,7 @@ def solve_heat1d(backend): device_data = backend.zeros(n) y = sun.N_VMake_Cuda(n, host_data, device_data, sunctx) - problem = Heat1DCudaProblem(backend, n=n) + problem = Heat1DCudaProblem(backend, n=n, k=k) problem.set_init_cond(y) stepper = ark.ARKStepCreate(problem.f, None, 0.0, y, sunctx) @@ -165,6 +195,11 @@ def solve_heat1d(backend): print(f" {t:10.6f} {rms:10.6f}") tout = min(tout + tf / nt, tf) + uexact = exact_semidiscrete_solution(n, k, tf) + max_error = np.max(np.abs(host_data - uexact)) + print(f"\nFinal max error vs exact semi-discrete solution = {max_error:.6e}") + np.testing.assert_allclose(host_data, uexact, rtol=1e-4, atol=1e-8) + return host_data.copy(), backend.to_numpy(device_data) From 99314494b0751404ad73d6a5a47a007bd4406ae2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Jul 2026 05:30:24 +0000 Subject: [PATCH 3/3] bindings: update litgen CUDA generated headers Signed-off-by: GitHub --- .../nvector/nvector_cuda_generated.hpp | 2 -- .../sunmemory/sunmemory_cuda_generated.hpp | 16 ++++++++++++---- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/bindings/sundials4py/nvector/nvector_cuda_generated.hpp b/bindings/sundials4py/nvector/nvector_cuda_generated.hpp index bdfde029f8..37ff8c8b6b 100644 --- a/bindings/sundials4py/nvector/nvector_cuda_generated.hpp +++ b/bindings/sundials4py/nvector/nvector_cuda_generated.hpp @@ -15,8 +15,6 @@ 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_VGetLength_Cuda", N_VGetLength_Cuda, nb::arg("x")); - m.def("N_VEnableFusedOps_Cuda", N_VEnableFusedOps_Cuda, nb::arg("v"), nb::arg("tf")); diff --git a/bindings/sundials4py/sunmemory/sunmemory_cuda_generated.hpp b/bindings/sundials4py/sunmemory/sunmemory_cuda_generated.hpp index 43c2a9a4c5..913e99e4a3 100644 --- a/bindings/sundials4py/sunmemory/sunmemory_cuda_generated.hpp +++ b/bindings/sundials4py/sunmemory/sunmemory_cuda_generated.hpp @@ -8,11 +8,19 @@ m.def( "SUNMemoryHelper_Cuda", [](SUNContext sunctx) -> std::shared_ptr> { - auto helper = SUNMemoryHelper_Cuda(sunctx); - return our_make_shared, - SUNMemoryHelperDeleter>(helper); + auto SUNMemoryHelper_Cuda_adapt_return_type_to_shared_ptr = + [](SUNContext sunctx) + -> std::shared_ptr> + { + auto lambda_result = SUNMemoryHelper_Cuda(sunctx); + + return our_make_shared, + SUNMemoryHelperDeleter>(lambda_result); + }; + + return SUNMemoryHelper_Cuda_adapt_return_type_to_shared_ptr(sunctx); }, - nb::arg("sunctx"), nb::keep_alive<0, 1>()); + nb::arg("sunctx"), "nb::keep_alive<0, 1>()", nb::keep_alive<0, 1>()); // #ifdef __cplusplus // // #endif