reduction

Author: leofang

cuda_core/build_hooks.py

@@ -28,8 +28,8 @@
 @functools.cache
 def _get_cuda_paths():
     CUDA_PATH = os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME", None))
-    if not CUDA_PATH:
-        raise RuntimeError("Environment variable CUDA_PATH or CUDA_HOME is not set")
+    if CUDA_PATH is None:
+        return None
     CUDA_PATH = CUDA_PATH.split(os.pathsep)
     print("CUDA paths:", CUDA_PATH, flush=True)
     return CUDA_PATH
@@ -46,7 +46,11 @@ def _get_cuda_version_from_cuda_h(cuda_home=None):
     Returns the integer (e.g. 13000) or None if not found / on error.
     """
     if cuda_home is None:
-        cuda_home = _get_cuda_paths()[0]
+        cuda_home = _get_cuda_paths()
+        if cuda_home is None:
+            return None
+        else:
+            cuda_home = cuda_home[0]
 
     cuda_h = pathlib.Path(cuda_home) / "include" / "cuda.h"
     if not cuda_h.is_file():
@@ -62,58 +66,32 @@ def _get_cuda_version_from_cuda_h(cuda_home=None):
     if not m:
         return None
     print(f"CUDA_VERSION from {cuda_h}:", m.group(1), flush=True)
-    return m.group(1)
-
+    return int(m.group(1))
 
-def _get_cuda_driver_version_linux():
-    """
-    Linux-only. Try to load `libcuda.so` via standard dynamic library lookup
-    and call `CUresult cuDriverGetVersion(int* driverVersion)`.
 
-    Returns:
-        int  : driver version (e.g., 12040 for 12.4), if successful.
-        None : on any failure (load error, missing symbol, non-success CUresult).
+def _get_cuda_driver_version():
     """
-    CUDA_SUCCESS = 0
-
-    libcuda_so = "libcuda.so.1"
-    cdll_mode = os.RTLD_NOW | os.RTLD_GLOBAL
-    try:
-        # Use system search paths only; do not provide an absolute path.
-        # Make symbols globally available to any dependent libraries.
-        lib = ctypes.CDLL(libcuda_so, mode=cdll_mode)
-    except OSError:
-        return None
+    Try to load ``libcuda.so`` or ``nvcuda.dll`` via standard dynamic library lookup
+    and call ``cuDriverGetVersion``.
 
-    # int cuDriverGetVersion(int* driverVersion);
-    lib.cuDriverGetVersion.restype = ctypes.c_int  # CUresult
-    lib.cuDriverGetVersion.argtypes = [ctypes.POINTER(ctypes.c_int)]
-
-    out = ctypes.c_int(0)
-    rc = lib.cuDriverGetVersion(ctypes.byref(out))
-    if rc != CUDA_SUCCESS:
-        return None
-
-    print(f"CUDA_VERSION from {libcuda_so}:", int(out.value), flush=True)
-    return int(out.value)
-
-
-def _get_cuda_driver_version_windows():
-    """
-    Windows-only. Load `nvcuda.dll` via normal system search and call
-    CUresult cuDriverGetVersion(int* driverVersion).
-
-    Returns:
-        int  : driver version (e.g., 12040 for 12.4), if successful.
-        None : on any failure (load error, missing symbol, non-success CUresult).
+    Returns the integer (e.g. 13000) or None if not found / on error.
     """
     CUDA_SUCCESS = 0
 
-    try:
-        # WinDLL => stdcall (CUDAAPI on Windows), matches CUDA Driver API.
-        lib = ctypes.WinDLL("nvcuda.dll")
-    except OSError:
-        return None
+    if sys.platform == "win32":
+        try:
+            # WinDLL => stdcall (CUDAAPI on Windows), matches CUDA Driver API.
+            lib = ctypes.WinDLL("nvcuda.dll")
+        except OSError:
+            return None
+    else:
+        cdll_mode = os.RTLD_NOW | os.RTLD_GLOBAL
+        try:
+            # Use system search paths only; do not provide an absolute path.
+            # Make symbols globally available to any dependent libraries.
+            lib = ctypes.CDLL("libcuda.so.1", mode=cdll_mode)
+        except OSError:
+            return None
 
     # int cuDriverGetVersion(int* driverVersion);
     cuDriverGetVersion = lib.cuDriverGetVersion
@@ -125,7 +103,7 @@ def _get_cuda_driver_version_windows():
     if rc != CUDA_SUCCESS:
         return None
 
-    print("CUDA_VERSION from nvcuda.dll:", int(out.value), flush=True)
+    print("CUDA_VERSION from driver:", int(out.value), flush=True)
     return int(out.value)
 
 
@@ -145,14 +123,11 @@ def _get_proper_cuda_bindings_major_version() -> str:
         return cuda_major
 
     cuda_version = _get_cuda_version_from_cuda_h()
-    if cuda_version and len(cuda_version) > 3:
-        return cuda_version[:-3]
+    if cuda_version:
+        return str(cuda_version // 1000)
 
     # also for local development
-    if sys.platform == "win32":
-        cuda_version = _get_cuda_driver_version_windows()
-    else:
-        cuda_version = _get_cuda_driver_version_linux()
+    cuda_version = _get_cuda_driver_version()
     if cuda_version:
         return str(cuda_version // 1000)