Updating CUDA extension and removing CUFFT callbacks

Author: ASKabalan

CMakeLists.txt

@@ -40,7 +40,7 @@ if(CMAKE_CUDA_COMPILER)
       OPTIONAL_COMPONENTS Development.SABIModule)
     execute_process(
       COMMAND "${Python_EXECUTABLE}" "-c"
-              "from jax.extend import ffi; print(ffi.include_dir())"
+              "from jax import ffi; print(ffi.include_dir())"
       OUTPUT_STRIP_TRAILING_WHITESPACE
       OUTPUT_VARIABLE XLA_DIR)
     message(STATUS "XLA include directory: ${XLA_DIR}")

lib/include/s2fft_callbacks.h

@@ -1,3 +1,14 @@
+/**
+ * @file s2fft_callbacks.h
+ * @brief CUDA CUFFT callbacks for HEALPix spherical harmonic transforms
+ * 
+ * @note CUFFT CALLBACKS DEPRECATED: This implementation no longer uses cuFFT callbacks.
+ * The previous callback-based approach has been replaced with direct kernel launches
+ * for better performance and maintainability. The files s2fft_callbacks.h and 
+ * s2fft_callbacks.cc are no longer used and can be considered orphaned.
+ */
+
+
 #ifndef _S2FFT_CALLBACKS_CUH_
 #define _S2FFT_CALLBACKS_CUH_
 

lib/include/s2fft_kernels.h

@@ -9,17 +9,77 @@
 #include <cufftXt.h>
 typedef long long int int64;
 
+/**
+ * @file s2fft_kernels.h
+ * @brief CUDA kernels for HEALPix spherical harmonic transforms
+ * 
+ * @note CUFT CALLBACKS DEPRECATED: This implementation no longer uses cuFFT callbacks.
+ * The previous callback-based approach has been replaced with direct kernel launches
+ * for better performance and maintainability. The files s2fft_callbacks.h and 
+ * s2fft_callbacks.cc are no longer used and can be considered orphaned.
+ */
+
 namespace s2fftKernels {
 
 enum fft_norm { FORWARD = 1, BACKWARD = 2, ORTHO = 3, NONE = 4 };
 
+/**
+ * @brief Launches the spectral folding CUDA kernel.
+ *
+ * This function configures and launches the spectral_folding kernel with
+ * appropriate grid and block dimensions. It performs spectral folding operations
+ * on ring-ordered data, transforming from Fourier coefficient space to HEALPix
+ * pixel space with optional FFT shifting.
+ *
+ * @tparam complex The complex type (cufftComplex or cufftDoubleComplex).
+ * @param data Input data array containing Fourier coefficients per ring.
+ * @param output Output array for folded HEALPix pixel data.
+ * @param nside The HEALPix Nside parameter.
+ * @param L The harmonic band limit.
+ * @param shift Flag indicating whether to apply FFT shifting.
+ * @param stream CUDA stream for kernel execution.
+ * @return HRESULT indicating success or failure.
+ */
 template <typename complex>
 HRESULT launch_spectral_folding(complex* data, complex* output, const int& nside, const int& L,
                                 const bool& shift, cudaStream_t stream);
+
+/**
+ * @brief Launches the spectral extension CUDA kernel.
+ *
+ * This function configures and launches the spectral_extension kernel with
+ * appropriate grid and block dimensions. It performs the inverse operation of
+ * spectral folding, extending HEALPix pixel data back to full Fourier coefficient
+ * space by mapping folded frequency components to their appropriate positions.
+ *
+ * @tparam complex The complex type (cufftComplex or cufftDoubleComplex).
+ * @param data Input array containing folded HEALPix pixel data.
+ * @param output Output array for extended Fourier coefficients per ring.
+ * @param nside The HEALPix Nside parameter.
+ * @param L The harmonic band limit.
+ * @param stream CUDA stream for kernel execution.
+ * @return HRESULT indicating success or failure.
+ */
 template <typename complex>
 HRESULT launch_spectral_extension(complex* data, complex* output, const int& nside, const int& L,
                                   cudaStream_t stream);
 
+/**
+ * @brief Launches the shift/normalize CUDA kernel for HEALPix data processing.
+ *
+ * This function configures and launches the shift_normalize_kernel with appropriate
+ * grid and block dimensions. It handles both single and double precision complex
+ * types and applies the requested normalization and shifting operations to HEALPix
+ * pixel data on a per-ring basis.
+ *
+ * @tparam complex The complex type (cufftComplex or cufftDoubleComplex).
+ * @param stream CUDA stream for kernel execution.
+ * @param data Input/output array of HEALPix pixel data (in-place processing).
+ * @param nside The HEALPix Nside parameter.
+ * @param apply_shift Flag indicating whether to apply FFT shifting.
+ * @param norm Normalization type (0=by nphi, 1=by sqrt(nphi), 2=no normalization).
+ * @return HRESULT indicating success or failure.
+ */
 template <typename complex>
 HRESULT launch_shift_normalize_kernel(cudaStream_t stream,
                                       complex* data,  // In-place data buffer