try to defer driver loading

Author: leofang

cuda_core/cuda/core/experimental/_launcher.py

@@ -16,6 +16,25 @@
 from cuda.core.experimental._utils import CUDAError, check_or_create_options, handle_return
 
 
+# TODO: revisit this treatment for py313t builds
+_inited = False
+_use_ex = None
+
+
+def _lazy_init():
+    global _inited
+    if _inited:
+        return
+
+    global _use_ex
+    # binding availability depends on cuda-python version
+    _py_major_minor = tuple(int(v) for v in (
+        importlib.metadata.version("cuda-python").split(".")[:2]))
+    _driver_ver = handle_return(cuda.cuDriverGetVersion())
+    _use_ex = (_driver_ver >= 11080) and (_py_major_minor >= (11, 8))
+    _inited = True
+
+
 @dataclass
 class LaunchConfig:
     """
@@ -41,6 +60,8 @@ def __post_init__(self):
         if self.shmem_size is None:
             self.shmem_size = 0
 
+        _lazy_init()
+
     def _cast_to_3_tuple(self, cfg):
         if isinstance(cfg, int):
             if cfg < 1:
@@ -65,13 +86,6 @@ def _cast_to_3_tuple(self, cfg):
             raise ValueError
 
 
-# binding availability depends on cuda-python version
-py_major_minor = tuple(int(v) for v in (
-    importlib.metadata.version("cuda-python").split(".")[:2]))
-driver_ver = handle_return(cuda.cuDriverGetVersion())
-use_ex = (driver_ver >= 11080) and (py_major_minor >= (11, 8))
-
-
 def launch(kernel, config, *kernel_args):
     if not isinstance(kernel, Kernel):
         raise ValueError
@@ -89,7 +103,7 @@ def launch(kernel, config, *kernel_args):
     # "new" module loading APIs are in use). We check both binding & driver versions here
     # mainly to see if the "Ex" API is available and if so we use it, as it's more feature
     # rich.
-    if use_ex:
+    if _use_ex:
         drv_cfg = cuda.CUlaunchConfig()
         drv_cfg.gridDimX, drv_cfg.gridDimY, drv_cfg.gridDimZ = config.grid
         drv_cfg.blockDimX, drv_cfg.blockDimY, drv_cfg.blockDimZ = config.block

cuda_core/cuda/core/experimental/_module.py

@@ -16,18 +16,33 @@
     },
 }
 
-# binding availability depends on cuda-python version
-py_major_ver = int(importlib.metadata.version("cuda-python").split(".")[0])
-if py_major_ver >= 12:
-    _backend["new"] = {
-        "file": cuda.cuLibraryLoadFromFile,
-        "data": cuda.cuLibraryLoadData,
-        "kernel": cuda.cuLibraryGetKernel,
-    }
-    _kernel_ctypes = (cuda.CUfunction, cuda.CUkernel)
-else:
-    _kernel_ctypes = (cuda.CUfunction,)
-driver_ver = handle_return(cuda.cuDriverGetVersion())
+
+# TODO: revisit this treatment for py313t builds
+_inited = False
+_py_major_ver = None
+_driver_ver = None
+_kernel_ctypes = None
+
+
+def _lazy_init():
+    global _inited
+    if _inited:
+        return
+
+    global _py_major_ver, _driver_ver, _kernel_ctypes
+    # binding availability depends on cuda-python version
+    _py_major_ver = int(importlib.metadata.version("cuda-python").split(".")[0])
+    if _py_major_ver >= 12:
+        _backend["new"] = {
+            "file": cuda.cuLibraryLoadFromFile,
+            "data": cuda.cuLibraryLoadData,
+            "kernel": cuda.cuLibraryGetKernel,
+        }
+        _kernel_ctypes = (cuda.CUfunction, cuda.CUkernel)
+    else:
+        _kernel_ctypes = (cuda.CUfunction,)
+    _driver_ver = handle_return(cuda.cuDriverGetVersion())
+    _inited = True
 
 
 class Kernel:
@@ -58,13 +73,14 @@ def __init__(self, module, code_type, jit_options=None, *,
                  symbol_mapping=None):
         if code_type not in self._supported_code_type:
             raise ValueError
+        _lazy_init()
         self._handle = None
 
-        backend = "new" if (py_major_ver >= 12 and driver_ver >= 12000) else "old"
+        backend = "new" if (_py_major_ver >= 12 and _driver_ver >= 12000) else "old"
         self._loader = _backend[backend]
 
         if isinstance(module, str):
-            if driver_ver < 12000 and jit_options is not None:
+            if _driver_ver < 12000 and jit_options is not None:
                 raise ValueError
             module = module.encode()
             self._handle = handle_return(self._loader["file"](module))