Update JAX Binding to use FFI

Author: ASKabalan

.pre-commit-config.yaml

@@ -4,3 +4,10 @@ repos:
     hooks:
       - id: ruff
       - id: ruff-format
+  - repo: https://github.com/pre-commit/mirrors-clang-format
+    rev: v18.1.4
+    hooks:
+      - id: clang-format
+        files: '\.(c|cc|cpp|h|hpp|cxx|hh|cu|cuh)$'
+        exclude: '^third_party/|/pybind11/'
+        name: clang-format

CMakeLists.txt

@@ -28,13 +28,15 @@ if(CMAKE_CUDA_COMPILER)
   else()
     find_package(CUDAToolkit REQUIRED)
 
-    find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
+    # Add the executable
+    find_package(Python 3.8
+    REQUIRED COMPONENTS Interpreter Development.Module
+    OPTIONAL_COMPONENTS Development.SABIModule)
+    set(XLA_DIR ${Python_SITELIB}/jaxlib/include)
+    message(STATUS "XLA_DIR: ${XLA_DIR}")
 
     # Detect the installed nanobind package and import it into CMake
-    execute_process(
-      COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
-      OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)
-    find_package(nanobind CONFIG REQUIRED) 
+    find_package(nanobind CONFIG REQUIRED)
 
     nanobind_add_module(_s2fft STABLE_ABI  
                     ${CMAKE_CURRENT_LIST_DIR}/lib/src/extensions.cc
@@ -45,7 +47,10 @@ if(CMAKE_CUDA_COMPILER)
                   )
 
     target_link_libraries(_s2fft PRIVATE CUDA::cudart_static CUDA::cufft_static CUDA::culibos)
-    target_include_directories(_s2fft PUBLIC ${CMAKE_CURRENT_LIST_DIR}/lib/include) 
+    target_include_directories(_s2fft PUBLIC 
+                                ${CMAKE_CURRENT_LIST_DIR}/lib/include 
+                                ${XLA_DIR}
+    ) 
     set_target_properties(_s2fft PROPERTIES 
                         LINKER_LANGUAGE CUDA
                         CUDA_SEPARABLE_COMPILATION ON)

lib/include/kernel_helpers.h

@@ -28,16 +28,16 @@ void s2fft_nphi_2_rings(float *data, int nside);
 
 class s2fftDescriptor {
 public:
-    int nside;
-    int harmonic_band_limit;
+    int64_t nside;
+    int64_t harmonic_band_limit;
     bool reality;
 
     bool forward = true;
     s2fftKernels::fft_norm norm = s2fftKernels::BACKWARD;
     bool shift = true;
     bool double_precision = false;
 
-    s2fftDescriptor(int nside, int harmonic_band_limit, bool reality, bool forward = true,
+    s2fftDescriptor(int64_t nside, int64_t harmonic_band_limit, bool reality, bool forward = true,
                     s2fftKernels::fft_norm norm = s2fftKernels::BACKWARD, bool shift = true,
                     bool double_precision = false)
             : nside(nside),
@@ -95,7 +95,7 @@ namespace std {
 template <>
 struct hash<s2fft::s2fftDescriptor> {
     std::size_t operator()(const s2fft::s2fftDescriptor &k) const {
-        size_t hash = std::hash<int>()(k.nside) ^ (std::hash<int>()(k.harmonic_band_limit) << 1) ^
+        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);
         return hash;