Add support for cutile jax ffi

Co-authored-by: yinuol <yinuol@nvidia.com>
Signed-off-by: Jay Gu <jagu@nvidia.com>

Author: haijieg

CMakeLists.txt

@@ -90,6 +90,9 @@ else()
     message(STATUS "Fetching dlpack")
 endif()
 
+include(cmake/FetchXLAHeaders.cmake)
+fetch_xla_headers()
+
 add_subdirectory(cext)
 
 

cext/CMakeLists.txt

@@ -36,8 +36,10 @@ endif()
 
 set(cext_include_dirs
     ${dlpack_INCLUDE_DIR}
+    ${xla_INCLUDE_DIR}
     ${Python_INCLUDE_DIRS}
-    ${CUDAToolkit_INCLUDE_DIRS})
+    ${CUDAToolkit_INCLUDE_DIRS}
+)
 
 
 # Build a static library first, so that we could reuse it for several build targets
@@ -49,6 +51,8 @@ add_library(_cext_static STATIC
     py.cpp
     stream_buffer.cpp
     tile_kernel.cpp
+    xla_ffi.cpp
+    xla_ffi_py.cpp
     )
 
 target_compile_options(_cext_static PUBLIC ${cext_compile_flags} ${nostdlib_flags})

cext/cuda_loader.cpp

@@ -77,16 +77,9 @@ static Status cuda_loader_init(DriverApi& driver_api) {
 
 static constexpr int MIN_DRIVER_VERSION = 13000;
 
-#ifdef Py_GIL_DISABLED
-static PyMutex g_driver_api_mutex = {0};
-#endif
-
 Result<const DriverApi*> get_driver_api() {
     static bool initialized;
     static DriverApi instance;
-#ifdef Py_GIL_DISABLED
-    PyCriticalSectionGuard guard(&g_driver_api_mutex);
-#endif
     if (!initialized) {
         if (!cuda_loader_init(instance))
             return ErrorRaised;

cext/cuda_loader.h

@@ -33,6 +33,7 @@
     X(cuEventDestroy, 2000) \
     X(cuEventQuery, 2000) \
     X(cuEventRecord, 2000) \
+    X(cuKernelGetFunction, 12000) \
     X(cuMemAlloc, 3020) \
     X(cuMemAllocHost, 3020) \
     X(cuMemFree, 3020) \

cext/hash.h

@@ -89,5 +89,3 @@ struct Hash<T, typename std::enable_if_t<std::is_enum_v<T>> > {
         return Hash<uint64_t>::hash(static_cast<uint64_t>(val), h);
     }
 };
-
-

cext/module.cpp

@@ -6,6 +6,7 @@
 
 #include "tile_kernel.h"
 #include "cuda_helper.h"
+#include "xla_ffi_py.h"
 
 #ifdef _WIN32
 extern "C" int _fltused = 0;
@@ -47,6 +48,10 @@ PyMODINIT_FUNC PyInit__cext() {
     if (!cuda_helper_init(m.get()))
         return nullptr;
 
+    if (!xla_ffi_init(m.get()))
+        return nullptr;
+
+
     return m.release();
 }
 

cext/py.h

@@ -277,6 +277,22 @@ static inline PyPtr try_import(const char* modname, SavedException* exc = nullpt
     return ret;
 }
 
+class GILGuard {
+public:
+    GILGuard(const GILGuard&) = delete;
+    void operator=(const GILGuard&) = delete;
+
+    GILGuard() {
+        gstate = PyGILState_Ensure();
+    }
+
+    ~GILGuard() {
+        PyGILState_Release(gstate);
+    }
+private:
+    PyGILState_STATE gstate;
+};
+
 #ifdef Py_GIL_DISABLED
 class PyCriticalSectionGuard {
     public: