ensure program/linker have consistent backend and handle

Author: leofang

cuda_core/cuda/core/experimental/_linker.py

@@ -484,9 +484,14 @@ def _input_type_from_code_type(self, code_type: str):
 
     @property
     def handle(self):
-        """Return the linker handle object."""
+        """Return the underlying handle object."""
         return self._mnff.handle
 
+    @property
+    def backend(self) -> str:
+        """Return this Linker instance's underlying backend."""
+        return "nvJitLink" if self._mnff.use_nvjitlink else "driver"
+
     def close(self):
         """Destroy this linker."""
         self._mnff.close()

cuda_core/cuda/core/experimental/_program.py

@@ -382,7 +382,7 @@ def __init__(self, code, code_type, options: ProgramOptions = None):
             # TODO: allow tuples once NVIDIA/cuda-python#72 is resolved
 
             self._mnff.handle = handle_return(nvrtc.nvrtcCreateProgram(code.encode(), b"", 0, [], []))
-            self._backend = "nvrtc"
+            self._backend = "NVRTC"
             self._linker = None
 
         elif code_type == "ptx":
@@ -391,7 +391,7 @@ def __init__(self, code, code_type, options: ProgramOptions = None):
             self._linker = Linker(
                 ObjectCode._init(code.encode(), code_type), options=self._translate_program_options(options)
             )
-            self._backend = "linker"
+            self._backend = self._linker.backend
         else:
             raise NotImplementedError
 
@@ -445,9 +445,9 @@ def compile(self, target_type, name_expressions=(), logs=None):
 
         """
         if target_type not in self._supported_target_type:
-            raise NotImplementedError
+            raise ValueError(f"the target type {target_type} is not supported")
 
-        if self._backend == "nvrtc":
+        if self._backend == "NVRTC":
             if target_type == "ptx" and not self._can_load_generated_ptx():
                 warn(
                     "The CUDA driver version is older than the backend version. "
@@ -489,15 +489,15 @@ def compile(self, target_type, name_expressions=(), logs=None):
 
             return ObjectCode._init(data, target_type, symbol_mapping=symbol_mapping)
 
-        if self._backend == "linker":
+        if self._backend in ("nvJitLink", "driver"):
             return self._linker.link(target_type)
 
     @property
-    def backend(self):
-        """Return the backend type string associated with this program."""
+    def backend(self) -> str:
+        """Return this Program instance's underlying backend."""
         return self._backend
 
     @property
     def handle(self):
-        """Return the program handle object."""
+        """Return the underlying handle object."""
         return self._mnff.handle

cuda_core/tests/test_linker.py

@@ -17,8 +17,8 @@
 device_function_b = "__device__ int B() { return 0; }"
 device_function_c = "__device__ int C(int a, int b) { return a + b; }"
 
-culink_backend = _linker._decide_nvjitlink_or_driver()
-if not culink_backend:
+is_culink_backend = _linker._decide_nvjitlink_or_driver()
+if not is_culink_backend:
     from cuda.bindings import nvjitlink
 
 
@@ -54,7 +54,7 @@ def compile_ltoir_functions(init_cuda):
     LinkerOptions(arch=ARCH, debug=True),
     LinkerOptions(arch=ARCH, lineinfo=True),
 ]
-if not culink_backend:
+if not is_culink_backend:
     options += [
         LinkerOptions(arch=ARCH, time=True),
         LinkerOptions(arch=ARCH, optimize_unused_variables=True),
@@ -85,24 +85,25 @@ def test_linker_init(compile_ptx_functions, options):
     linker = Linker(*compile_ptx_functions, options=options)
     object_code = linker.link("cubin")
     assert isinstance(object_code, ObjectCode)
+    assert linker.backend == ("driver" if is_culink_backend else "nvJitLink")
 
 
 def test_linker_init_invalid_arch(compile_ptx_functions):
-    err = AttributeError if culink_backend else nvjitlink.nvJitLinkError
+    err = AttributeError if is_culink_backend else nvjitlink.nvJitLinkError
     with pytest.raises(err):
         options = LinkerOptions(arch="99", ptx=True)
         Linker(*compile_ptx_functions, options=options)
 
 
-@pytest.mark.skipif(culink_backend, reason="culink does not support ptx option")
+@pytest.mark.skipif(is_culink_backend, reason="culink does not support ptx option")
 def test_linker_link_ptx_nvjitlink(compile_ltoir_functions):
     options = LinkerOptions(arch=ARCH, link_time_optimization=True, ptx=True)
     linker = Linker(*compile_ltoir_functions, options=options)
     linked_code = linker.link("ptx")
     assert isinstance(linked_code, ObjectCode)
 
 
-@pytest.mark.skipif(not culink_backend, reason="nvjitlink requires lto for ptx linking")
+@pytest.mark.skipif(not is_culink_backend, reason="nvjitlink requires lto for ptx linking")
 def test_linker_link_ptx_culink(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
     linker = Linker(*compile_ptx_functions, options=options)

cuda_core/tests/test_program.py

@@ -14,6 +14,8 @@
 from cuda.core.experimental._module import Kernel, ObjectCode
 from cuda.core.experimental._program import Program, ProgramOptions
 
+is_culink_backend = _linker._decide_nvjitlink_or_driver()
+
 
 @pytest.fixture(scope="module")
 def ptx_code_object():
@@ -50,7 +52,7 @@ def ptx_code_object():
 def test_cpp_program_with_various_options(init_cuda, options):
     code = 'extern "C" __global__ void my_kernel() {}'
     program = Program(code, "c++", options)
-    assert program.backend == "nvrtc"
+    assert program.backend == "NVRTC"
     program.compile("ptx")
     program.close()
     assert program.handle is None
@@ -65,8 +67,7 @@ def test_cpp_program_with_various_options(init_cuda, options):
     ProgramOptions(prec_sqrt=True),
     ProgramOptions(fma=True),
 ]
-if not _linker._decide_nvjitlink_or_driver():
-    print("Using nvjitlink as the backend because decide() returned false")
+if not is_culink_backend:
     options += [
         ProgramOptions(time=True),
         ProgramOptions(split_compile=True),
@@ -76,7 +77,7 @@ def test_cpp_program_with_various_options(init_cuda, options):
 @pytest.mark.parametrize("options", options)
 def test_ptx_program_with_various_options(init_cuda, ptx_code_object, options):
     program = Program(ptx_code_object._module.decode(), "ptx", options=options)
-    assert program.backend == "linker"
+    assert program.backend == ("driver" if is_culink_backend else "nvJitLink")
     program.compile("cubin")
     program.close()
     assert program.handle is None
@@ -85,7 +86,7 @@ def test_ptx_program_with_various_options(init_cuda, ptx_code_object, options):
 def test_program_init_valid_code_type():
     code = 'extern "C" __global__ void my_kernel() {}'
     program = Program(code, "c++")
-    assert program.backend == "nvrtc"
+    assert program.backend == "NVRTC"
     assert program.handle is not None
 
 
@@ -125,14 +126,14 @@ def test_program_compile_valid_target_type(init_cuda):
 def test_program_compile_invalid_target_type():
     code = 'extern "C" __global__ void my_kernel() {}'
     program = Program(code, "c++")
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError):
         program.compile("invalid_target")
 
 
 def test_program_backend_property():
     code = 'extern "C" __global__ void my_kernel() {}'
     program = Program(code, "c++")
-    assert program.backend == "nvrtc"
+    assert program.backend == "NVRTC"
 
 
 def test_program_handle_property():