begin replacing pynvjitlinker

Author: brandon-b-miller

numba_cuda/numba/cuda/cudadrv/driver.py

@@ -47,6 +47,9 @@
 except ImportError:
     NvJitLinker, NvJitLinkError = None, None
 
+from cuda.bindings.nvjitlink import nvJitLinkError
+from cuda.bindings import nvjitlink as _nvjitlink
+
 USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING
 
 if USE_NV_BINDING:
@@ -2594,7 +2597,7 @@ def new(cls,
                 "Enabling pynvjitlink requires CUDA 12."
             )
         if config.CUDA_ENABLE_PYNVJITLINK:
-            linker = PyNvJitLinker
+            linker = PyNvJitLinkerNew
 
         elif config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY:
             linker = MVCLinker
@@ -3024,6 +3027,143 @@ def complete(self):
         return bytes(np.ctypeslib.as_array(cubin_ptr, shape=(size,)))
 
 
+class PyNvJitLinkerNew(Linker):
+    def __init__(
+        self,
+        max_registers=None,
+        lineinfo=False,
+        cc=None,
+        lto=False,
+        additional_flags=None,
+    ):
+
+        if cc is None:
+            raise RuntimeError("NvJitLink requires CC to be specified")
+        if not any(isinstance(cc, t) for t in [list, tuple]):
+            raise TypeError("`cc` must be a list or tuple of length 2")
+
+        sm_ver = f"{cc[0] * 10 + cc[1]}"
+        arch = f"-arch=sm_{sm_ver}"
+        options = [arch]
+        if max_registers:
+            options.append(f"-maxrregcount={max_registers}")
+        if lineinfo:
+            options.append("-lineinfo")
+        if lto:
+            options.append("-lto")
+        if additional_flags is not None:
+            options.extend(additional_flags)
+
+        self.handle = _nvjitlink.create(len(options), options)
+        self.lto = lto
+        self.options = options
+
+        self._info_log = None
+        self._error_log = None
+        self._complete = False
+
+    @property
+    def info_log(self):
+        return self._info_log
+
+    @property
+    def error_log(self):
+        return self._error_log
+
+    def add_data(self, input_type, data, name):
+        if self._complete:
+            raise nvJitLinkError("Cannot add data to already-completeted link")
+
+        try:
+            _nvjitlink.add_data(
+                self.handle,
+                input_type.value,
+                data,
+                len(data),
+                name
+            )
+        except RuntimeError as e:
+            log_size = _nvjitlink.get_error_log_size(self.handle)
+            log = bytearray(log_size)
+            self._info_log = _nvjitlink.get_error_log(self.handle, log)
+
+            log_size = _nvjitlink.get_info_log_size(self.handle)
+            log = bytearray(log_size)
+            self._error_log = _nvjitlink.get_info_log(self.handle, log)
+            raise nvJitLinkError(f"{e}\n{self.error_log}")
+
+    def add_cubin(self, cubin, name=None):
+        name = name or "unnamed-cubin"
+        self.add_data(_nvjitlink.InputType.CUBIN, cubin, name)
+
+    def add_ptx(self, ptx, name=None):
+        name = name or "unnamed-ptx"
+        self.add_data(_nvjitlink.InputType.PTX, ptx, name)
+
+    def add_ltoir(self, ltoir, name=None):
+        name = name or "unnamed-ltoir"
+        self.add_data(_nvjitlink.InputType.LTOIR, ltoir, name)
+
+    def add_object(self, object_, name=None):
+        name = name or "unnamed-object"
+        self.add_data(_nvjitlink.InputType.OBJECT, object_, name)
+
+    def add_fatbin(self, fatbin, name=None):
+        name = name or "unnamed-fatbin"
+        self.add_data(_nvjitlink.InputType.FATBIN, fatbin, name)
+
+    def add_library(self, library, name=None):
+        self.add_data(_nvjitlink.InputType.LIBRARY, library, name)
+
+    def add_file(self, path, kind):
+        try:
+            with open(path, "rb") as f:
+                data = f.read()
+        except FileNotFoundError:
+            raise LinkerError(f"{path} not found")
+
+        name = pathlib.Path(path).name
+        self.add_data(data, kind, name)
+
+    def get_linked_cubin(self):
+        try:
+            _nvjitlink.complete(self.handle)
+            self._complete = True
+            size = _nvjitlink.get_linked_cubin_size(self.handle)
+            cubin = bytearray(size)
+
+            _nvjitlink.get_linked_cubin(self.handle, cubin)
+            return cubin
+
+        except nvJitLinkError as e:
+            size = _nvjitlink.get_error_log_size(self.handle)
+            log = bytearray(size)
+            self._error_log = _nvjitlink.get_error_log(self.handle, log)
+            raise nvJitLinkError(f"{e}\n{self.error_log}")
+        finally:
+            size = _nvjitlink.get_info_log_size(self.handle)
+            log = bytearray(size)
+            self._info_log = _nvjitlink.get_info_log(self.handle, log)
+
+    def get_linked_ptx(self):
+        try:
+            _nvjitlink.complete(self.handle)
+            self._complete = True
+            size = _nvjitlink.get_linked_ptx_size(self.handle)
+            return _nvjitlink.get_linked_ptx(self.handle, size)
+        except RuntimeError as e:
+            self._error_log = _nvjitlink.get_error_log(self.handle)
+            raise NvJitLinkError(f"{e}\n{self.error_log}")
+        finally:
+            self._info_log = _nvjitlink.get_info_log(self.handle)
+
+    def complete(self):
+        try:
+            return self.get_linked_cubin()
+        except NvJitLinkError as e:
+            raise LinkerError from e
+
+
 class PyNvJitLinker(Linker):
     def __init__(
         self,