Delay cuda driver load until launch

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

Author: haijieg

cext/CMakeLists.txt

@@ -74,6 +74,7 @@ target_link_options(_cext_shared PUBLIC ${cext_link_flags} ${nostdlib_flags} -Wl
 add_executable(test_stream_buffer
     test/test_stream_buffer.cpp
     cuda_loader.cpp
+    cuda_helper.cpp
     memory.cpp
     )
 target_compile_options(test_stream_buffer PUBLIC ${cext_compile_flags} ${test_coverage_options})

cext/cuda_helper.cpp

@@ -6,40 +6,52 @@
 #include "cuda_loader.h"
 
 
-const char* get_cuda_error(CUresult res) {
+const char* get_cuda_error(const DriverApi* driver, CUresult res) {
     const char* str = nullptr;
-    g_cuGetErrorString(res, &str);
+    driver->cuGetErrorString(res, &str);
     return str ? str : "Unknown error";
 }
 
-void try_init_cuda() {
-    ErrorGuard guard;
-    CUresult res = g_cuInit(0);
+Status check_driver_version(const DriverApi* driver, int minimum_version) {
+    int version;
+    CUresult res = driver->cuDriverGetVersion(&version);
     if (res != CUDA_SUCCESS) {
-        raise(PyExc_RuntimeError, "cuInit: %s", get_cuda_error(res));
-        SavedException exc = save_raised_exception();
-        LOG_PYTHON_ERROR("warning", exc, "Failed to initialized CUDA");
+        PyErr_Format(PyExc_RuntimeError, "cuDriverGetVersion: %s", get_cuda_error(driver, res));
+        return ErrorRaised;
+    }
+    if (version < minimum_version) {
+        int major = version / 1000;
+        int minor = (version % 1000) / 10;
+        int required_major = minimum_version / 1000;
+        PyErr_Format(PyExc_RuntimeError,
+                     "Minimum driver version required is %d.0, got %d.%d",
+                     required_major, major, minor);
+        return ErrorRaised;
     }
+    return OK;
 }
 
 PyObject* get_max_grid_size(PyObject *self, PyObject *args) {
     int device_id;
     if (!PyArg_ParseTuple(args, "i", &device_id))
         return NULL;
 
+    Result<const DriverApi*> driver = get_driver_api();
+    if (!driver.is_ok()) return NULL;
+
     CUdevice dev;
-    CUresult res = g_cuDeviceGet(&dev, device_id);
+    CUresult res = (*driver)->cuDeviceGet(&dev, device_id);
     if (res != CUDA_SUCCESS)
-        return PyErr_Format(PyExc_RuntimeError, "cuDeviceGet: %s", get_cuda_error(res));
+        return PyErr_Format(PyExc_RuntimeError, "cuDeviceGet: %s", get_cuda_error(*driver, res));
 
     int max_grid_size[3];
     for (int i = 0; i < 3; ++i) {
-        res = g_cuDeviceGetAttribute(&max_grid_size[i],
+        res = (*driver)->cuDeviceGetAttribute(&max_grid_size[i],
             static_cast<CUdevice_attribute>(CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X + i),
             dev);
         if (res != CUDA_SUCCESS) {
             return PyErr_Format(PyExc_RuntimeError,
-                                "cuDeviceGetAttribute: %s", get_cuda_error(res));
+                                "cuDeviceGetAttribute: %s", get_cuda_error(*driver, res));
         }
     }
     return Py_BuildValue("(iii)", max_grid_size[0], max_grid_size[1], max_grid_size[2]);
@@ -48,26 +60,36 @@ PyObject* get_max_grid_size(PyObject *self, PyObject *args) {
 PyObject* get_compute_capability(PyObject *self, PyObject *Py_UNUSED(ignored)) {
     int major, minor;
     CUdevice dev;
-    CUresult res = g_cuDeviceGet(&dev, 0);
+
+    Result<const DriverApi*> driver_result = get_driver_api();
+    if (!driver_result.is_ok()) return NULL;
+    const DriverApi* d = *driver_result;
+
+    CUresult res = d->cuDeviceGet(&dev, 0);
     if (res != CUDA_SUCCESS) {
-        return PyErr_Format(PyExc_RuntimeError, "cuDeviceGet: %s", get_cuda_error(res));
+        return PyErr_Format(PyExc_RuntimeError, "cuDeviceGet: %s", get_cuda_error(d, res));
     }
-    res = g_cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev);
+    res = d->cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev);
     if (res != CUDA_SUCCESS) {
-        return PyErr_Format(PyExc_RuntimeError, "cuDeviceGetAttribute: %s", get_cuda_error(res));
+        return PyErr_Format(PyExc_RuntimeError, "cuDeviceGetAttribute: %s", get_cuda_error(d, res));
     }
-    res = g_cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev);
+    res = d->cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev);
     if (res != CUDA_SUCCESS) {
-        return PyErr_Format(PyExc_RuntimeError, "cuDeviceGetAttribute: %s", get_cuda_error(res));
+        return PyErr_Format(PyExc_RuntimeError, "cuDeviceGetAttribute: %s", get_cuda_error(d, res));
     }
     return Py_BuildValue("(ii)", major, minor);
 }
 
 PyObject* get_driver_version(PyObject *self, PyObject *Py_UNUSED(ignored)) {
     int major, minor;
-    CUresult res = g_cuDriverGetVersion(&major);
+
+    Result<const DriverApi*> driver_result = get_driver_api();
+    if (!driver_result.is_ok()) return NULL;
+    const DriverApi* d = *driver_result;
+
+    CUresult res = d->cuDriverGetVersion(&major);
     if (res != CUDA_SUCCESS) {
-        return PyErr_Format(PyExc_RuntimeError, "cuDriverGetVersion: %s", get_cuda_error(res));
+        return PyErr_Format(PyExc_RuntimeError, "cuDriverGetVersion: %s", get_cuda_error(d, res));
     }
     minor = (major % 1000) / 10;
     major = major / 1000;

cext/cuda_helper.h

@@ -9,8 +9,12 @@
 #include "py.h"
 #include <cuda.h>
 
+struct DriverApi;
+
 Status cuda_helper_init(PyObject* m);
 
-const char* get_cuda_error(CUresult res);
+const char* get_cuda_error(const DriverApi*, CUresult res);
+
+void try_cuInit(const DriverApi*);
 
-void try_init_cuda();
+Status check_driver_version(const DriverApi*, int minimum_version);

cext/cuda_loader.cpp

@@ -50,11 +50,12 @@ F get_proc_address(cuGetProcAddress_v2_t getter,
 FOREACH_CUDA_FUNCTION_TO_LOAD(DEFINE_CUDA_FUNCTION_GLOBAL)
 
 #define GET_PROC_ADDRESS(name, cuda_ver) \
-        if (!(g_##name = get_proc_address<decltype(name)*>(_cuGetProcAddress, #name, cuda_ver))) \
+        if (!(driver_api.name = \
+                    get_proc_address<decltype(name)*>(_cuGetProcAddress, #name, cuda_ver))) \
             return ErrorRaised;
 
 
-Status cuda_loader_init() {
+static Status cuda_loader_init(DriverApi& driver_api) {
     PyPtr load_libcuda_mod = steal(PyImport_ImportModule("cuda.tile._load_libcuda"));
     if (!load_libcuda_mod) return ErrorRaised;
 
@@ -72,3 +73,22 @@ Status cuda_loader_init() {
 
     return OK;
 }
+
+
+static constexpr int MIN_DRIVER_VERSION = 13000;
+
+Result<const DriverApi*> get_driver_api() {
+    static bool initialized;
+    static DriverApi instance;
+    if (!initialized) {
+        if (!cuda_loader_init(instance))
+            return ErrorRaised;
+        CUresult res = instance.cuInit(0);
+        if (res != CUDA_SUCCESS)
+            return raise(PyExc_RuntimeError, "cuInit: %s", get_cuda_error(&instance, res));
+        if (!check_driver_version(&instance, MIN_DRIVER_VERSION))
+            return ErrorRaised;
+        initialized = true;
+    }
+    return &instance;
+}

cext/cuda_loader.h

@@ -9,8 +9,6 @@
 #include "py.h"
 #include <cuda.h>
 
-Status cuda_loader_init();
-
 #define FOREACH_CUDA_FUNCTION_TO_LOAD(X) \
     X(cuInit, 2000) \
     X(cuLibraryLoadFromFile, 12000) \
@@ -47,8 +45,10 @@ Status cuda_loader_init();
 
 
 #define DECLARE_CUDA_FUNC_EXTERN(name, _cuda_version) \
-    extern decltype(name)* g_##name;
-
-FOREACH_CUDA_FUNCTION_TO_LOAD(DECLARE_CUDA_FUNC_EXTERN)
+    decltype(::name)* name;
 
+struct DriverApi {
+    FOREACH_CUDA_FUNCTION_TO_LOAD(DECLARE_CUDA_FUNC_EXTERN)
+};
 
+Result<const DriverApi*> get_driver_api();

cext/module.cpp

@@ -4,7 +4,6 @@
 
 #include "py.h"
 
-#include "cuda_loader.h"
 #include "tile_kernel.h"
 #include "cuda_helper.h"
 
@@ -20,11 +19,6 @@ static PyModuleDef module_def = {
 };
 
 PyMODINIT_FUNC PyInit__cext() {
-    if (!cuda_loader_init())
-        return nullptr;
-
-    try_init_cuda();
-
     PyPtr m = steal(PyModule_Create(&module_def));
     if (!m) return nullptr;