Fix memory illegal access issue

Add comprehensive documentation and fix dependency issues for CUDA FFT integration.

This commit introduces extensive docstrings and inline comments across the C++ and Python codebase, particularly for the CUDA FFT implementation. It also addresses a dependency issue in  to ensure proper installation and functionality.

Key changes include:
- no more CUDA Malloc .. all memory is allocated in Python by XLA
- Added detailed docstrings to C++ header files
- Enhanced inline comments in C++ source files to explain complex logic and algorithms.
- Updated to relax JAX version dependency, resolving installation issues.
- Refined docstrings and comments in Python files for clarity and consistency.
- Cleaned up debug print statements

Author: ASKabalan

lib/include/plan_cache.h

@@ -1,4 +1,3 @@
-
 #ifndef PLAN_CACHE_H
 #define PLAN_CACHE_H
 
@@ -9,36 +8,84 @@
 #include "hresult.h"
 #include "s2fft.h"
 #include <unordered_map>
+#include <type_traits>
 
 namespace s2fft {
 
+/**
+ * @brief Manages and caches s2fftExec instances to optimize resource usage.
+ *
+ * This class implements the singleton pattern to ensure only one instance
+ * of the PlanCache exists throughout the application. It stores pre-initialized
+ * s2fftExec objects based on their descriptors (parameters like nside, L, etc.)
+ * to avoid redundant initialization, which can be computationally expensive.
+ */
 class PlanCache {
 public:
+    /**
+     * @brief Returns the singleton instance of the PlanCache.
+     *
+     * @return A reference to the single PlanCache instance.
+     */
     static PlanCache &GetInstance() {
         static PlanCache instance;
         return instance;
     }
 
-    HRESULT GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptr<s2fftExec<cufftComplex>> &executor);
+    /**
+     * @brief Retrieves an s2fftExec instance from the cache or initializes a new one.
+     *
+     * This templated method attempts to find an existing s2fftExec instance
+     * matching the provided descriptor in its internal cache (m_Descriptors32 or m_Descriptors64)
+     * based on the Complex type T. If a matching instance is found, it is returned.
+     * Otherwise, a new s2fftExec instance is created, initialized with the descriptor,
+     * and then stored in the cache before being returned.
+     *
+     * @tparam T The complex type (cufftComplex or cufftDoubleComplex) of the s2fftExec instance.
+     * @param descriptor The s2fftDescriptor containing the parameters for the FFT.
+     * @param executor A shared_ptr that will point to the retrieved or newly initialized s2fftExec instance.
+     * @return HRESULT indicating success (S_OK if new, S_FALSE if from cache) or failure.
+     */
+    template <typename T>
+    HRESULT GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptr<s2fftExec<T>> &executor);
 
-    HRESULT GetS2FFTExec(s2fftDescriptor &descriptor,
-                         std::shared_ptr<s2fftExec<cufftDoubleComplex>> &executor);
+    /**
+     * @brief Clears all cached s2fftExec instances.
+     *
+     * This method is typically called during application shutdown to release
+     * all resources held by the cached FFT plans.
+     */
+    void Finalize();
 
-    ~PlanCache() {}
+    /**
+     * @brief Destructor for PlanCache.
+     *
+     * Ensures that Finalize() is called when the PlanCache instance is destroyed,
+     * performing necessary cleanup.
+     */
+    ~PlanCache();
 
 private:
     bool is_initialized = false;
 
+    // Unordered maps to store cached s2fftExec instances for double and single precision
     std::unordered_map<s2fftDescriptor, std::shared_ptr<s2fftExec<cufftDoubleComplex>>,
                        std::hash<s2fftDescriptor>, std::equal_to<>>
             m_Descriptors64;
     std::unordered_map<s2fftDescriptor, std::shared_ptr<s2fftExec<cufftComplex>>, std::hash<s2fftDescriptor>,
                        std::equal_to<>>
             m_Descriptors32;
 
+    /**
+     * @brief Private constructor for PlanCache.
+     *
+     * Initializes the PlanCache instance. This constructor is private to enforce
+     * the singleton pattern.
+     */
     PlanCache();
 
 public:
+    // Delete copy constructor and assignment operator to prevent copying
     PlanCache(PlanCache const &) = delete;
     void operator=(PlanCache const &) = delete;
 };

lib/include/s2fft.h

@@ -1,4 +1,3 @@
-
 #ifndef S2FFT_H
 #define S2FFT_H
 
@@ -19,13 +18,49 @@
 
 namespace s2fft {
 
+/**
+ * @brief Returns the appropriate cuFFT C2C type for a given complex type.
+ *
+ * This function is overloaded for `cufftDoubleComplex` and `cufftComplex`
+ * to return `CUFFT_Z2Z` (double precision) or `CUFFT_C2C` (single precision)
+ * respectively.
+ *
+ * @param dummy A dummy complex object used for type deduction.
+ * @return The corresponding cuFFT C2C type.
+ */
 static cufftType get_cufft_type_c2c(cufftDoubleComplex) { return CUFFT_Z2Z; }
 static cufftType get_cufft_type_c2c(cufftComplex) { return CUFFT_C2C; }
 
+/**
+ * @brief Transforms data from ring-based indexing to nphi-based indexing.
+ *
+ * This function is a placeholder for the actual implementation which would
+ * reorder data in memory according to the specified indexing scheme.
+ *
+ * @param data Pointer to the input/output data.
+ * @param nside The HEALPix Nside parameter.
+ */
 void s2fft_rings_2_nphi(float *data, int nside);
 
+/**
+ * @brief Transforms data from nphi-based indexing to ring-based indexing.
+ *
+ * This function is a placeholder for the actual implementation which would
+ * reorder data in memory according to the specified indexing scheme.
+ *
+ * @param data Pointer to the input/output data.
+ * @param nside The HEALPix Nside parameter.
+ */
 void s2fft_nphi_2_rings(float *data, int nside);
 
+/**
+ * @brief Descriptor class for s2fft operations.
+ *
+ * This class encapsulates all the necessary parameters to define a unique
+ * Spherical Harmonic Transform (SHT) operation, including Nside, harmonic
+ * band limit, reality, adjoint flag, forward/backward transform direction,
+ * normalization, shifting, and double precision usage.
+ */
 class s2fftDescriptor {
 public:
     int64_t nside;
@@ -38,6 +73,18 @@ class s2fftDescriptor {
     bool shift = true;
     bool double_precision = false;
 
+    /**
+     * @brief Constructs an s2fftDescriptor object.
+     *
+     * @param nside The HEALPix Nside parameter.
+     * @param harmonic_band_limit The harmonic band limit L.
+     * @param reality Flag indicating if the signal is real.
+     * @param adjoint Flag indicating if the adjoint transform is to be performed.
+     * @param forward Flag indicating if it's a forward transform (default: true).
+     * @param norm The FFT normalization type (default: BACKWARD).
+     * @param shift Flag indicating if FFT shifting should be applied (default: true).
+     * @param double_precision Flag indicating if double precision should be used (default: false).
+     */
     s2fftDescriptor(int64_t nside, int64_t harmonic_band_limit, bool reality, bool adjoint,
                     bool forward = true, s2fftKernels::fft_norm norm = s2fftKernels::BACKWARD,
                     bool shift = true, bool double_precision = false)
@@ -50,53 +97,130 @@ class s2fftDescriptor {
               shift(shift),
               double_precision(double_precision) {}
 
+    /**
+     * @brief Default constructor for s2fftDescriptor.
+     */
     s2fftDescriptor() = default;
+
+    /**
+     * @brief Destructor for s2fftDescriptor.
+     */
     ~s2fftDescriptor() = default;
 
+    /**
+     * @brief Equality operator for s2fftDescriptor.
+     *
+     * Compares two s2fftDescriptor objects for equality based on their member values.
+     *
+     * @param other The other s2fftDescriptor to compare against.
+     * @return True if the descriptors are equal, false otherwise.
+     */
     bool operator==(const s2fftDescriptor &other) const {
         return nside == other.nside && harmonic_band_limit == other.harmonic_band_limit &&
                reality == other.reality && norm == other.norm && shift == other.shift &&
                double_precision == other.double_precision;
     }
 };
 
+/**
+ * @brief Executes Spherical Harmonic Transform (SHT) operations.
+ *
+ * This templated class provides methods for initializing FFT plans and executing
+ * forward and backward SHTs. It manages cuFFT handles and internal offsets
+ * required for the transforms.
+ *
+ * @tparam Complex The complex type (cufftComplex or cufftDoubleComplex) for the FFT operations.
+ */
 template <typename Complex>
 class s2fftExec {
-    friend class PlanCache;
+    friend class PlanCache;  // Allows PlanCache to access private members for caching
 
 public:
+    /**
+     * @brief Default constructor for s2fftExec.
+     */
     s2fftExec() {}
-    ~s2fftExec() {}
-
-    HRESULT Initialize(const s2fftDescriptor &descriptor, size_t &worksize);
 
-    HRESULT Forward(const s2fftDescriptor &desc, cudaStream_t stream, Complex *data);
+    /**
+     * @brief Destructor for s2fftExec.
+     */
+    ~s2fftExec() {}
 
-    HRESULT Backward(const s2fftDescriptor &desc, cudaStream_t stream, Complex *data);
+    /**
+     * @brief Initializes the FFT plans for the SHT.
+     *
+     * This method sets up the necessary cuFFT plans for both polar and equatorial
+     * rings based on the provided descriptor. It also calculates and stores the
+     * maximum required workspace size (m_work_size).
+     *
+     * @param descriptor The s2fftDescriptor containing the parameters for the FFT.
+     * @return HRESULT indicating success or failure.
+     */
+    HRESULT Initialize(const s2fftDescriptor &descriptor);
+
+    /**
+     * @brief Executes the forward Spherical Harmonic Transform.
+     *
+     * This method performs the forward FFT operations on the input data
+     * across polar and equatorial rings using the pre-initialized cuFFT plans.
+     *
+     * @param desc The s2fftDescriptor for the current transform.
+     * @param stream The CUDA stream to use for execution.
+     * @param data Pointer to the input/output data on the device.
+     * @param workspace Pointer to the workspace memory on the device.
+     * @param callback_params Pointer to device memory containing callback parameters.
+     * @return HRESULT indicating success or failure.
+     */
+    HRESULT Forward(const s2fftDescriptor &desc, cudaStream_t stream, Complex *data, Complex *workspace,
+                    int64 *callback_params);
+
+    /**
+     * @brief Executes the backward Spherical Harmonic Transform.
+     *
+     * This method performs the inverse FFT operations on the input data
+     * across polar and equatorial rings using the pre-initialized cuFFT plans.
+     *
+     * @param desc The s2fftDescriptor for the current transform.
+     * @param stream The CUDA stream to use for execution.
+     * @param data Pointer to the input/output data on the device.
+     * @param workspace Pointer to the workspace memory on the device.
+     * @param callback_params Pointer to device memory containing callback parameters.
+     * @return HRESULT indicating success or failure.
+     */
+    HRESULT Backward(const s2fftDescriptor &desc, cudaStream_t stream, Complex *data, Complex *workspace,
+                     int64 *callback_params);
 
 public:
+    // cuFFT handles for polar and equatorial FFT plans
     std::vector<cufftHandle> m_polar_plans;
     cufftHandle m_equator_plan;
     std::vector<cufftHandle> m_inverse_polar_plans;
     cufftHandle m_inverse_equator_plan;
+
+    // Parameters defining the SHT geometry and data layout
     int m_nside;
     int m_equatorial_ring_num;
     int64 m_total_pixels;
     int64 m_equatorial_offset_start;
     int64 m_equatorial_offset_end;
     std::vector<int64> m_upper_ring_offsets;
     std::vector<int64> m_lower_ring_offsets;
-
-    // Callback params stored for cleanup purposes
-    // thrust::device_vector<cb_params> m_cb_params;
+    size_t m_work_size = 0;  // Maximum workspace size required for FFT plans
 };
 
 }  // namespace s2fft
 
 namespace std {
+/**
+ * @brief Custom hash specialization for s2fftDescriptor.
+ *
+ * This specialization allows s2fftDescriptor objects to be used as keys
+ * in `std::unordered_map` by providing a hash function.
+ */
 template <>
 struct hash<s2fft::s2fftDescriptor> {
     std::size_t operator()(const s2fft::s2fftDescriptor &k) const {
+        // Combine hash values of individual members
         size_t hash = std::hash<int64_t>()(k.nside) ^ (std::hash<int64_t>()(k.harmonic_band_limit) << 1) ^
                       (std::hash<bool>()(k.reality) << 2) ^ (std::hash<int>()(k.norm) << 3) ^
                       (std::hash<bool>()(k.shift) << 4) ^ (std::hash<bool>()(k.double_precision) << 5);

lib/include/s2fft_callbacks.h

@@ -12,10 +12,44 @@
 typedef long long int int64;
 
 namespace s2fftKernels {
+/**
+ * @brief Defines the normalization types for FFT operations.
+ */
 enum fft_norm { FORWARD = 1, BACKWARD = 2, ORTHO = 3, NONE = 4 };
 
-HRESULT setCallback(cufftHandle forwardPlan, cufftHandle backwardPlan, int64 *params_dev, bool shift,
-                    bool equator, bool doublePrecision, fft_norm norm);
+/**
+ * @brief Sets cuFFT callbacks specifically for a forward FFT plan.
+ *
+ * This function configures the cuFFT library to use custom callbacks
+ * for normalization and shifting operations during forward FFT execution.
+ *
+ * @param plan The cuFFT handle for the forward FFT plan.
+ * @param params_dev Pointer to device memory containing parameters for the callbacks.
+ * @param shift Boolean flag indicating whether to apply FFT shifting.
+ * @param equator Boolean flag indicating if the current operation is for the equatorial ring.
+ * @param doublePrecision Boolean flag indicating if double precision is used.
+ * @param norm The FFT normalization type to apply.
+ * @return HRESULT indicating success or failure.
+ */
+HRESULT setForwardCallback(cufftHandle plan, int64 *params_dev, bool shift, bool equator,
+                           bool doublePrecision, fft_norm norm);
+
+/**
+ * @brief Sets cuFFT callbacks specifically for a backward FFT plan.
+ *
+ * This function configures the cuFFT library to use custom callbacks
+ * for normalization and shifting operations during backward FFT execution.
+ *
+ * @param plan The cuFFT handle for the inverse FFT plan.
+ * @param params_dev Pointer to device memory containing parameters for the callbacks.
+ * @param shift Boolean flag indicating whether to apply FFT shifting.
+ * @param equator Boolean flag indicating if the current operation is for the equatorial ring.
+ * @param doublePrecision Boolean flag indicating if double precision is used.
+ * @param norm The FFT normalization type to apply.
+ * @return HRESULT indicating success or failure.
+ */
+HRESULT setBackwardCallback(cufftHandle plan, int64 *params_dev, bool shift, bool equator,
+                            bool doublePrecision, fft_norm norm);
 }  // namespace s2fftKernels
 
 #endif