reuse backend decision logic in tests + some nitpicks

leofang · leofang · commit 677bd6df0158 · 2024-12-07T02:38:20.000Z
diff --git a/cuda_core/cuda/core/experimental/_linker.py b/cuda_core/cuda/core/experimental/_linker.py
@@ -20,29 +20,43 @@
 _nvjitlink_input_types = None  # populated if nvJitLink cannot be used
 
 
-def _lazy_init():
-    global _inited
-    if _inited:
+# Note: this function is reused in the tests
+def _decide_nvjitlink_or_driver():
+    """Returns True if falling back to the cuLink* driver APIs."""
+    global _driver_ver, _driver, _nvjitlink
+    if _driver or _nvjitlink:
         return
 
-    global _driver, _driver_input_types, _driver_ver, _nvjitlink, _nvjitlink_input_types
     _driver_ver = handle_return(cuda.cuDriverGetVersion())
     _driver_ver = (_driver_ver // 1000, (_driver_ver % 1000) // 10)
     try:
-        from cuda.bindings import nvjitlink
+        from cuda.bindings import nvjitlink as _nvjitlink
         from cuda.bindings._internal import nvjitlink as inner_nvjitlink
     except ImportError:
         # binding is not available
-        nvjitlink = None
+        _nvjitlink = None
     else:
         if inner_nvjitlink._inspect_function_pointer("__nvJitLinkVersion") == 0:
             # binding is available, but nvJitLink is not installed
-            nvjitlink = None
-        elif _driver_ver > nvjitlink.version():
+            _nvjitlink = None
+
+    if _nvjitlink is None:
+        _driver = cuda
+        return True
+    else:
+        return False
+
+
+def _lazy_init():
+    global _inited, _nvjitlink_input_types, _driver_input_types
+    if _inited:
+        return
+
+    _decide_nvjitlink_or_driver()
+    if _nvjitlink:
+        if _driver_ver > _nvjitlink.version():
             # TODO: nvJitLink is not new enough, warn?
             pass
-    if nvjitlink:
-        _nvjitlink = nvjitlink
         _nvjitlink_input_types = {
             "ptx": _nvjitlink.InputType.PTX,
             "cubin": _nvjitlink.InputType.CUBIN,
@@ -51,8 +65,6 @@ def _lazy_init():
             "object": _nvjitlink.InputType.OBJECT,
         }
     else:
-        from cuda import cuda as _driver
-
         _driver_input_types = {
             "ptx": _driver.CUjitInputType.CU_JIT_INPUT_PTX,
             "cubin": _driver.CUjitInputType.CU_JIT_INPUT_CUBIN,
diff --git a/cuda_core/tests/test_linker.py b/cuda_core/tests/test_linker.py
@@ -1,48 +1,37 @@
 import pytest
 
-from cuda.core.experimental import Linker, LinkerOptions, Program
+from cuda.core.experimental import Linker, LinkerOptions, Program, _linker
 from cuda.core.experimental._module import ObjectCode
 
 ARCH = "sm_80"  # use sm_80 for testing the oop nvJitLink wrapper
 
-
-device_function_a = """
-__device__ int B();
-__device__ int C(int a, int b);
+kernel_a = """
+extern __device__ int B();
+extern __device__ int C(int a, int b);
 __global__ void A() { int result = C(B(), 1);}
 """
 device_function_b = "__device__ int B() { return 0; }"
 device_function_c = "__device__ int C(int a, int b) { return a + b; }"
 
-culink_backend = False
-try:
-    from cuda.bindings import nvjitlink  # noqa F401
-    from cuda.bindings._internal import nvjitlink as inner_nvjitlink
-except ImportError:
-    # binding is not available
-    culink_backend = True
-else:
-    if inner_nvjitlink._inspect_function_pointer("__nvJitLinkVersion") == 0:
-        # binding is available, but nvJitLink is not installed
-        culink_backend = True
+culink_backend = _linker._decide_nvjitlink_or_driver()
 
 
 @pytest.fixture(scope="function")
 def compile_ptx_functions(init_cuda):
-    # Without rdc (relocatable device code) option, the generated ptx will not included any unreferenced
+    # Without -rdc (relocatable device code) option, the generated ptx will not included any unreferenced
     # device functions, causing the link to fail
+    object_code_a_ptx = Program(kernel_a, "c++").compile("ptx", options=("-rdc=true",))
     object_code_b_ptx = Program(device_function_b, "c++").compile("ptx", options=("-rdc=true",))
     object_code_c_ptx = Program(device_function_c, "c++").compile("ptx", options=("-rdc=true",))
-    object_code_a_ptx = Program(device_function_a, "c++").compile("ptx", options=("-rdc=true",))
 
     return object_code_a_ptx, object_code_b_ptx, object_code_c_ptx
 
 
 @pytest.fixture(scope="function")
 def compile_ltoir_functions(init_cuda):
+    object_code_a_ltoir = Program(kernel_a, "c++").compile("ltoir", options=("-dlto",))
     object_code_b_ltoir = Program(device_function_b, "c++").compile("ltoir", options=("-dlto",))
     object_code_c_ltoir = Program(device_function_c, "c++").compile("ltoir", options=("-dlto",))
-    object_code_a_ltoir = Program(device_function_a, "c++").compile("ltoir", options=("-dlto",))
 
     return object_code_a_ltoir, object_code_b_ltoir, object_code_c_ltoir
 
