Add CUDA backend support for the heterobatch interface

CMakeLists.txt

@@ -15,6 +15,8 @@ find_package(nanobind CONFIG REQUIRED)
 
 option(NIFTY_LS_OPENMP "Enable OpenMP support in the compiled extensions" ON)
 option(NIFTY_LS_HETEROBATCH "Enable heterobatch support with finufft" ON)
+option(NIFTY_LS_CUDA "Enable CUDA backend (cuda_helper)" OFF)
+set(NIFTY_LS_CUFINUFFT_ROOT "" CACHE PATH "Prefix where cufinufft is installed (include/, lib/)")
 option(NIFTY_LS_SANITIZERS "Enable sanitizers for debugging" OFF)
 option(NIFTY_LS_COMPILE_WARNING_AS_ERROR "Treat compiler warnings as errors" OFF)
 
@@ -37,6 +39,26 @@ endif()
 
 set(NIFTY_LS_MODULES cpu_helpers chi2_helpers)
 
+if(NIFTY_LS_CUDA)
+    enable_language(CUDA)
+    find_package(CUDAToolkit REQUIRED)
+
+    # cufinufft C API (required for cuda_helper)
+    find_library(
+        CUFINUFFT_LIBRARY cufinufft
+        HINTS ${NIFTY_LS_CUFINUFFT_ROOT} ${NIFTY_LS_CUFINUFFT_ROOT}/lib ${NIFTY_LS_CUFINUFFT_ROOT}/lib64
+    )
+    find_path(
+        CUFINUFFT_INCLUDE_DIR cufinufft.h
+        HINTS ${NIFTY_LS_CUFINUFFT_ROOT} ${NIFTY_LS_CUFINUFFT_ROOT}/include
+    )
+    if(NOT CUFINUFFT_LIBRARY OR NOT CUFINUFFT_INCLUDE_DIR)
+        message(FATAL_ERROR "cufinufft library and headers are required when NIFTY_LS_CUDA=ON")
+    endif()
+
+    list(APPEND NIFTY_LS_MODULES cuda_helper chi2_cuda_helper)
+endif()
+
 if(NIFTY_LS_HETEROBATCH)
     # Download CPM.cmake if not already available
     if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/cmake/CPM.cmake)
@@ -64,7 +86,15 @@ endif()
 
 
 foreach(MODULE_NAME ${NIFTY_LS_MODULES})
-    nanobind_add_module(${MODULE_NAME} src/nifty_ls/${MODULE_NAME}.cpp NOMINSIZE STABLE_ABI FREE_THREADED)
+    set(MODULE_SRC "src/nifty_ls/${MODULE_NAME}.cpp")
+    if(${MODULE_NAME} STREQUAL "cuda_helper")
+        set(MODULE_SRC "src/nifty_ls/${MODULE_NAME}.cu")
+    endif()
+    if(${MODULE_NAME} STREQUAL "chi2_cuda_helper")
+        set(MODULE_SRC "src/nifty_ls/${MODULE_NAME}.cu")
+    endif()
+
+    nanobind_add_module(${MODULE_NAME} ${MODULE_SRC} NOMINSIZE STABLE_ABI FREE_THREADED)
 
     target_compile_options(${MODULE_NAME} PRIVATE
         -Wall -Wextra -std=c++17
@@ -95,5 +125,26 @@ foreach(MODULE_NAME ${NIFTY_LS_MODULES})
         target_link_libraries(${MODULE_NAME} PRIVATE finufft)
     endif()
 
+    if(${MODULE_NAME} STREQUAL "cuda_helper" OR ${MODULE_NAME} STREQUAL "chi2_cuda_helper")
+        # cufinufft.h may live in ${prefix}/include or ${prefix}/include/cufinufft;
+        # finufft_common headers sit one level up, so add the parent include dir too.
+        get_filename_component(CUFINUFFT_INCLUDE_PARENT ${CUFINUFFT_INCLUDE_DIR} DIRECTORY)
+        target_include_directories(
+            ${MODULE_NAME}
+            PRIVATE
+            ${CUFINUFFT_INCLUDE_DIR}
+            ${CUFINUFFT_INCLUDE_PARENT}
+            ${CUFINUFFT_INCLUDE_PARENT}/finufft_common
+            ${CUFINUFFT_INCLUDE_PARENT}/../include
+        )
+        target_link_libraries(${MODULE_NAME} PRIVATE ${CUFINUFFT_LIBRARY} CUDA::cudart CUDA::cufft CUDA::cublas)
+        set_target_properties(
+            ${MODULE_NAME}
+            PROPERTIES
+                CUDA_SEPARABLE_COMPILATION ON
+                CUDA_ARCHITECTURES native
+        )
+    endif()
+
     install(TARGETS ${MODULE_NAME} LIBRARY DESTINATION nifty_ls)
 endforeach()

README.md

@@ -320,6 +320,29 @@ Individual periodograms in the heterobatch interface can have uniform batching;
 
 A similar dynamic dispatch effect can be achieved with [free-threaded Python](#free-threaded-parallelism), but as the dispatch is still done in Python, it will likely not be as performant as heterobatch.
 
+The `heterobatch` interface also provides a CUDA backend, which dispatches individual periodograms to cuFINUFFT at the C++ level using an OpenMP thread pool and CUDA streams. As a result, it can utilize GPU compute resources with very high efficiency.
+
+The CUDA backend requires cuFINUFFT as a prerequisite. Users must build cuFINUFFT separately and point to the installation path during configuration. For example:
+```bash
+### Build cufinufft
+git clone https://github.com/flatironinstitute/finufft.git
+cd finufft
+mkdir build && cd build
+cmake .. \
+  -DFINUFFT_USE_CUDA=ON \
+  -DFINUFFT_BUILD_TESTS=ON \
+  -DCMAKE_INSTALL_PREFIX=/custom/path/to/cufinufft
+
+cmake --build . -- -j
+ctest
+cmake --install .
+
+### Build nifty-ls with CUDA backend
+pip install -e .[cuda,dev] --no-build-isolation -Ccmake.args="-DNIFTY_LS_CUDA=ON -DNIFTY_LS_CUFINUFFT_ROOT=/custom/path/to/cufinufft"
+```
+
+CUDA compute sanitizers can be enabled by setting `-DNIFTY_LS_SANITIZERS=ON`, which helps detect CUDA memory errors during debugging.
+
 ### Free-Threaded Parallelism
 nifty-ls supports [free-threaded Python](https://docs.python.org/3/howto/free-threading-python.html)
 since version 1.1.0. With a free-threaded build of Python, efficient parallelism over many time
@@ -418,6 +441,8 @@ The 200× advantage of the GPU extends to even smaller $N$ in this case, since w
 
 We see that both multi-threaded finufft and cufinufft particularly benefit from batched transforms, as this exposes more parallelism and amortizes fixed latencies.
 
+With kernel fusion and CUDA streams, the `heterobatch` interface significantly reduces kernel launch overhead, which is particularly beneficial for small problem sizes. To better match GPU architectures, users can tune `TPB` (threads per block) to adjust GPU occupancy, and `CHUNK_F` to control memory usage.
+
 <!-- FUTURE: move to a readthedocs page and include the following performance plot
 
 In contrast, the following shows "batch mode" performance for the chi-squared method under the same setting and `nterms=4`:

scripts/bench.py

@@ -27,7 +27,9 @@
 DEFAULT_DTYPE = 'f8'
 DEFAULT_METHODS = [
     'cufinufft',
+    'cufinufft_CUDA',
     'cufinufft_chi2',
+    'cufinufft_chi2_CUDA',
     'finufft',
     'astropy',
     'finufft_par',
@@ -66,12 +68,24 @@ def do_nifty_cufinufft_chi2(*args, nterms=4, **kwargs):
     )
 
 
+def do_nifty_cufinufft_chi2_cuda(*args, nterms=4, **kwargs):
+    return nifty_ls.cufinufft_chi2_CUDA.lombscargle(
+        *args, **kwargs, nterms=nterms, cufinufft_kwargs={'eps': DEFAULT_EPS}
+    )
+
+
 def do_nifty_cufinufft(*args, **kwargs):
     return nifty_ls.cufinufft.lombscargle(
         *args, **kwargs, cufinufft_kwargs={'eps': DEFAULT_EPS}
     )
 
 
+def do_nifty_cufinufft_cuda(*args, **kwargs):
+    return nifty_ls.cufinufft_CUDA.lombscargle(
+        *args, **kwargs, cufinufft_kwargs={'eps': DEFAULT_EPS}
+    )
+
+
 def do_astropy_fast(t, y, dy, fmin, df, Nf, **astropy_kwargs):
     f0 = fmin
     y = np.atleast_2d(y)
@@ -139,7 +153,9 @@ def do_winding(t, y, dy, fmin, df, Nf, center_data=True, fit_mean=True, **kwargs
         *args, **kwargs, nthreads=1
     ),
     'cufinufft': do_nifty_cufinufft,
+    'cufinufft_CUDA': do_nifty_cufinufft_cuda,
     'cufinufft_chi2': do_nifty_cufinufft_chi2,
+    'cufinufft_chi2_CUDA': do_nifty_cufinufft_chi2_cuda,
     'astropy': do_astropy_fast,
     'astropy_brute': lambda *args, **kwargs: do_astropy_fast(
         *args, **kwargs, use_fft=False
@@ -234,6 +250,10 @@ def get_plot_kwargs(method, nthread_max=NTHREAD_MAX):
         label = 'nifty-ls (cufinufft)'
         color = 'C2'
         ls = '-'
+    elif method == 'cufinufft_CUDA':
+        label = 'nifty-ls (cufinufft CUDA)'
+        color = 'C2'
+        ls = '--'
     # Group 3: astropy and winding methods - different color, same family
     elif method == 'astropy':
         label = r'Astropy (${\tt fast}$ method)'

scripts/compare.py

@@ -13,6 +13,7 @@
 DEFAULT_LOGN = 3
 DEFAULT_METHODS = [
     'cufinufft',
+    'cufinufft_CUDA',
     'cufinufft_chi2',
     'finufft',
     'astropy',

src/nifty_ls/backends.py

@@ -15,9 +15,17 @@
     'HETEROBATCH_BACKEND_NAMES',
 ]
 
-CHI2_BACKEND_NAMES = ['finufft_chi2', 'cufinufft_chi2']
-STANDARD_BACKEND_NAMES = ['finufft', 'cufinufft']
-BACKEND_TYPE = Literal['auto', 'finufft', 'finufft_chi2', 'cufinufft', 'cufinufft_chi2']
+CHI2_BACKEND_NAMES = ['finufft_chi2', 'cufinufft_chi2', 'cufinufft_chi2_CUDA']
+STANDARD_BACKEND_NAMES = ['finufft', 'cufinufft', 'cufinufft_CUDA']
+BACKEND_TYPE = Literal[
+    'auto',
+    'finufft',
+    'finufft_chi2',
+    'cufinufft',
+    'cufinufft_chi2',
+    'cufinufft_chi2_CUDA',
+    'cufinufft_CUDA',
+]
 BACKEND_NAMES = list(get_args(BACKEND_TYPE))
 
 HETEROBATCH_STANDARD_BACKEND_NAMES = ['finufft_heterobatch']