Support wheels for Windows

Author: vzhurba01

cuda_bindings/cuda/bindings/_bindings/cynvrtc.pyx.in

@@ -6,8 +6,10 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 {{if 'Windows' == platform.system()}}
-import win32api
+import os
+import site
 import struct
+import win32api
 from pywintypes import error
 {{else}}
 cimport cuda.bindings._lib.dlfcn as dlfcn
@@ -44,11 +46,39 @@ cdef int cuPythonInit() except -1 nogil:
 
     # Load library
     {{if 'Windows' == platform.system()}}
-    LOAD_LIBRARY_SAFE_CURRENT_DIRS = 0x00002000
+    handle = NULL
     with gil:
+        # First check if the DLL has been loaded by 3rd parties
         try:
-            handle = win32api.LoadLibraryEx("nvrtc64_120_0.dll", 0, LOAD_LIBRARY_SAFE_CURRENT_DIRS)
+            handle = win32api.GetModuleHandle("nvrtc64_120_0.dll")
         except:
+            pass
+
+        # Try default search
+        if handle == NULL:
+            LOAD_LIBRARY_SAFE_CURRENT_DIRS = 0x00002000
+            try:
+                handle = win32api.LoadLibraryEx("nvrtc64_120_0.dll", 0, LOAD_LIBRARY_SAFE_CURRENT_DIRS)
+            except:
+                pass
+
+        # Check if DLLs are found within pip installations
+        if handle == NULL:
+            site_packages = [site.getusersitepackages()] + site.getsitepackages()
+            for sp in site_packages:
+                mod_path = os.path.join(sp, "nvidia", "cuda_nvrtc", "bin")
+                if not os.path.isdir(mod_path):
+                    continue
+                os.add_dll_directory(mod_path)
+            try:
+                handle = win32api.LoadLibraryEx(
+                    # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
+                    os.path.join(mod_path, "nvrtc64_120_0.dll"),
+                    0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
+            except:
+                pass
+
+        if handle == NULL:
             raise RuntimeError('Failed to LoadLibraryEx nvrtc64_120_0.dll')
     {{else}}
     handle = NULL

cuda_bindings/docs/source/release/12.x.y-notes.md

@@ -3,4 +3,22 @@
 Released on MM DD, 20YY.
 
 ## Highlights
-- Added bindings for nvJitLink. It requires nvJitLink from CUDA 12.3 or above.
+- Add bindings for nvJitLink. It requires nvJitLink from CUDA 12.3 or above.
+- Add optional dependencies to wheels for NVRTC and nvJitLink
+- Enable discovery and loading of shared library dependencies from wheels
+
+## Wheels support for optional dependencies
+
+Optional dependencies are added for packages:
+
+- nvidia-nvjitlink-cuXX
+- nvidia-cuda-nvrtc-cuXX
+
+Installing these dependencies with cuda-python can be done using:
+```{code-block} shell
+pip install cuda-python[all]
+```
+
+## Discovery and loading of shared library dependencies from wheels
+
+Shared library search paths for wheel builds are now extended to check site-packages. This allows users to seamlessly use their wheel installation of the CUDA Toolkit with cuda-python.

cuda_bindings/pyproject.toml

@@ -32,6 +32,12 @@ dependencies = [
   "pywin32; sys_platform == 'win32'",
 ]
 
+[project.optional-dependencies]
+all = [
+    "nvidia-cuda-nvrtc-cu12",
+    "nvidia-nvjitlink-cu12>=12.3"
+]
+
 [project.urls]
 Repository = "https://github.com/NVIDIA/cuda-python"
 Documentation = "https://nvidia.github.io/cuda-python/"

cuda_bindings/setup.py

@@ -301,19 +301,21 @@ def initialize_options(self):
             self.parallel = nthreads
 
     def build_extension(self, ext):
-        if building_wheel:
+        if building_wheel and sys.platform == "linux":
             # Strip binaries to remove debug symbols
             extra_linker_flags = ["-Wl,--strip-all"]
 
             # Allow extensions to discover libraries at runtime
             # relative their wheels installation.
-            ldflag = "-Wl,--disable-new-dtags"
             if ext.name == "cuda.bindings._bindings.cynvrtc":
-                ldflag += f",-rpath,$ORIGIN/../../../nvidia/cuda_nvrtc/lib"
+                ldflag = f"-Wl,--disable-new-dtags,-rpath,$ORIGIN/../../../nvidia/cuda_nvrtc/lib"
             elif ext.name == "cuda.bindings._internal.nvjitlink":
-                ldflag += f",-rpath,$ORIGIN/../../../nvidia/nvjitlink/lib"
+                ldflag = f"-Wl,--disable-new-dtags,-rpath,$ORIGIN/../../../nvidia/nvjitlink/lib"
+            else:
+                ldflag = None
 
-            extra_linker_flags.append(ldflag)
+            if ldflag:
+                extra_linker_flags.append(ldflag)
         else:
             extra_linker_flags = []