Avoid Using Device Compute Capability for Linker Class (#429)

When debugging with @gmarkall, we discovered that the linker class is
using the current device's compute capability for cuda source
compilation. Because `compile_ptx` can be used
to compile target code for a different architecture from current device.
This PR pass through the CC used to construct the linker class to the
cuda source compilation. This avoids
CC inconsistencies between compile and linking.

---------

Co-authored-by: Michael Wang <[email protected]>

Author: isVoid

numba_cuda/numba/cuda/cudadrv/driver.py

@@ -2836,10 +2836,7 @@ def add_ptx(self, ptx, name):
     def add_cu(self, cu, name):
         """Add CUDA source in a string to the link. The name of the source
         file should be specified in `name`."""
-        with driver.get_active_context() as ac:
-            dev = driver.get_device(ac.devnum)
-            cc = dev.compute_capability
-        ptx, log = nvrtc.compile(cu, name, cc)
+        ptx, log = nvrtc.compile(cu, name, self.cc)
 
         if config.DUMP_ASSEMBLY:
             print(("ASSEMBLY %s" % name).center(80, "-"))
@@ -3003,10 +3000,7 @@ def add_ptx(self, ptx, name="<cudapy-ptx>"):
         self._object_codes.append(obj)
 
     def add_cu(self, cu, name="<cudapy-cu>"):
-        with driver.get_active_context() as ac:
-            dev = driver.get_device(ac.devnum)
-            cc = dev.compute_capability
-        obj, log = nvrtc.compile(cu, name, cc, ltoir=self.lto)
+        obj, log = nvrtc.compile(cu, name, self.cc, ltoir=self.lto)
 
         if not self.lto and config.DUMP_ASSEMBLY:
             print(("ASSEMBLY %s" % name).center(80, "-"))
@@ -3117,6 +3111,7 @@ def __init__(self, max_registers=0, lineinfo=False, cc=None):
         if lineinfo:
             options[enums.CU_JIT_GENERATE_LINE_INFO] = c_void_p(1)
 
+        self.cc = cc
         if cc is None:
             # No option value is needed, but we need something as a placeholder
             options[enums.CU_JIT_TARGET_FROM_CUCONTEXT] = 1

numba_cuda/numba/cuda/tests/cudadrv/test_linker.py

@@ -5,7 +5,11 @@
 import warnings
 from numba import config
 from numba.cuda.testing import unittest
-from numba.cuda.testing import skip_on_cudasim, skip_if_cuda_includes_missing
+from numba.cuda.testing import (
+    skip_on_cudasim,
+    skip_if_cuda_includes_missing,
+    skip_if_nvjitlink_missing,
+)
 from numba.cuda.testing import CUDATestCase, test_data_dir
 from numba.cuda.cudadrv.driver import CudaAPIError, _Linker, LinkerError
 from numba.cuda import require_context
@@ -329,6 +333,18 @@ def test_get_local_mem_per_specialized(self):
         calc_size = np.dtype(np.float64).itemsize * LMEM_SIZE
         self.assertGreaterEqual(local_mem_size, calc_size)
 
+    @skip_if_nvjitlink_missing("nvJitLink not installed or new enough (>12.3)")
+    def test_link_for_different_cc(self):
+        linker = _Linker.new(cc=(7, 5), lto=True)
+        code = """
+__device__ int foo(int x) {
+    return x + 1;
+}
+"""
+        linker.add_cu(code, "foo")
+        ptx = linker.get_linked_ptx().decode()
+        assert "target sm_75" in ptx
+
 
 if __name__ == "__main__":
     unittest.main()