Fix bug in previous commit: need to break out of loop of suffix if site-package DLL was found. Using the most obvious approach to solve this problem: return immediately on success.

Author: rwgk

cuda_bindings/cuda/bindings/_internal/nvjitlink_windows.pyx

@@ -56,40 +56,30 @@ cdef load_library(const int driver_ver):
 
         # First check if the DLL has been loaded by 3rd parties
         try:
-            handle = win32api.GetModuleHandle(dll_name)
+            return win32api.GetModuleHandle(dll_name)
         except:
             pass
-        else:
-            break
 
         # Next, check if DLLs are installed via pip
         for sp in get_site_packages():
             mod_path = os.path.join(sp, "nvidia", "nvJitLink", "bin")
             if os.path.isdir(mod_path):
                 os.add_dll_directory(mod_path)
                 try:
-                    handle = win32api.LoadLibraryEx(
+                    return win32api.LoadLibraryEx(
                         # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
                         os.path.join(mod_path, dll_name),
                         0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
                 except:
                     pass
-                else:
-                    break
-        else:
-            # Finally, try default search
-            # Only reached if DLL wasn't found in any site-package path
-            try:
-                handle = win32api.LoadLibrary(dll_name)
-            except:
-                pass
-            else:
-                break
-    else:
-        raise RuntimeError('Failed to load nvJitLink')
-
-    assert handle != 0
-    return handle
+        # Finally, try default search
+        # Only reached if DLL wasn't found in any site-package path
+        try:
+            return win32api.LoadLibrary(dll_name)
+        except:
+            pass
+
+    raise RuntimeError('Failed to load nvJitLink')
 
 
 cdef int _check_or_init_nvjitlink() except -1 nogil:

cuda_bindings/cuda/bindings/_internal/nvvm_windows.pyx

@@ -54,11 +54,9 @@ cdef load_library(const int driver_ver):
 
         # First check if the DLL has been loaded by 3rd parties
         try:
-            handle = win32api.GetModuleHandle(dll_name)
+            return win32api.GetModuleHandle(dll_name)
         except:
             pass
-        else:
-            break
 
         # Next, check if DLLs are installed via pip or conda
         for sp in get_site_packages():
@@ -73,28 +71,21 @@ cdef load_library(const int driver_ver):
             if os.path.isdir(mod_path):
                 os.add_dll_directory(mod_path)
                 try:
-                    handle = win32api.LoadLibraryEx(
+                    return win32api.LoadLibraryEx(
                         # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
                         os.path.join(mod_path, dll_name),
                         0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
                 except:
                     pass
-                else:
-                    break
-        else:
-            # Finally, try default search
-            # Only reached if DLL wasn't found in any site-package path
-            try:
-                handle = win32api.LoadLibrary(dll_name)
-            except:
-                pass
-            else:
-                break
-    else:
-        raise RuntimeError('Failed to load nvvm')
 
-    assert handle != 0
-    return handle
+        # Finally, try default search
+        # Only reached if DLL wasn't found in any site-package path
+        try:
+            handle = win32api.LoadLibrary(dll_name)
+        except:
+            pass
+
+    raise RuntimeError('Failed to load nvvm')
 
 
 cdef int _check_or_init_nvvm() except -1 nogil: