diff --git a/numba_cuda/numba/cuda/cudadrv/driver.py b/numba_cuda/numba/cuda/cudadrv/driver.py index 0748931..fe46630 100644 --- a/numba_cuda/numba/cuda/cudadrv/driver.py +++ b/numba_cuda/numba/cuda/cudadrv/driver.py @@ -2667,19 +2667,38 @@ def add_cu_file(self, path): cu = f.read() self.add_cu(cu, os.path.basename(path)) - def add_file_guess_ext(self, path): + def add_file_guess_ext(self, path_or_code): """Add a file to the link, guessing its type from its extension.""" - ext = os.path.splitext(path)[1][1:] - if ext == '': - raise RuntimeError("Don't know how to link file with no extension") - elif ext == 'cu': - self.add_cu_file(path) - else: - kind = FILE_EXTENSION_MAP.get(ext, None) - if kind is None: - raise RuntimeError("Don't know how to link file with extension " - f".{ext}") - self.add_file(path, kind) + if isinstance(path_or_code, str): + ext = pathlib.Path(path_or_code).suffix + if ext == '': + raise RuntimeError( + "Don't know how to link file with no extension" + ) + elif ext == '.cu': + self.add_cu_file(path_or_code) + else: + kind = FILE_EXTENSION_MAP.get(ext, None) + if kind is None: + raise RuntimeError( + "Don't know how to link file with extension " + f".{ext}" + ) + self.add_file(path_or_code, kind) + return + else: + # Otherwise, we should have been given a LinkableCode object + if not isinstance(path_or_code, LinkableCode): + raise TypeError( + "Expected path to file or a LinkableCode object" + ) + + if path_or_code.kind == "cu": + self.add_cu(path_or_code.data, path_or_code.name) + else: + self.add_data( + path_or_code.data, path_or_code.kind, path_or_code.name + ) @abstractmethod def complete(self): @@ -3017,36 +3036,6 @@ def add_ltoir(self, ltoir, name=""): def add_object(self, obj, name=""): self._linker.add_object(obj, name) - def add_file_guess_ext(self, path_or_code): - # Numba's add_file_guess_ext expects to always be passed a path to a - # file that it will load from the filesystem to link. We augment it - # here with the ability to provide a file from memory. - - # To maintain compatibility with the original interface, all strings - # are treated as paths in the filesystem. - if isinstance(path_or_code, str): - # Upstream numba does not yet recognize LTOIR, so handle that - # separately here. - extension = pathlib.Path(path_or_code).suffix - if extension == ".ltoir": - self.add_file(path_or_code, "ltoir") - else: - # Use Numba's logic for non-LTOIR - super().add_file_guess_ext(path_or_code) - - return - - # Otherwise, we should have been given a LinkableCode object - if not isinstance(path_or_code, LinkableCode): - raise TypeError("Expected path to file or a LinkableCode object") - - if path_or_code.kind == "cu": - self.add_cu(path_or_code.data, path_or_code.name) - else: - self.add_data( - path_or_code.data, path_or_code.kind, path_or_code.name - ) - def add_file(self, path, kind): try: with open(path, "rb") as f: diff --git a/numba_cuda/numba/cuda/cudadrv/enums.py b/numba_cuda/numba/cuda/cudadrv/enums.py index 3431cf7..917dbb1 100644 --- a/numba_cuda/numba/cuda/cudadrv/enums.py +++ b/numba_cuda/numba/cuda/cudadrv/enums.py @@ -309,7 +309,10 @@ # Applicable options: PTX compiler options, ::CU_JIT_FALLBACK_STRATEGY CU_JIT_INPUT_LIBRARY = 4 -CU_JIT_NUM_INPUT_TYPES = 6 +# LTO IR +CU_JIT_INPUT_LTO_IR = 5 + +CU_JIT_NUM_INPUT_TYPES = 7 # Online compiler and linker options diff --git a/numba_cuda/numba/cuda/cudadrv/mappings.py b/numba_cuda/numba/cuda/cudadrv/mappings.py index 3750324..15fe79b 100644 --- a/numba_cuda/numba/cuda/cudadrv/mappings.py +++ b/numba_cuda/numba/cuda/cudadrv/mappings.py @@ -19,4 +19,5 @@ 'lib': enums.CU_JIT_INPUT_LIBRARY, 'cubin': enums.CU_JIT_INPUT_CUBIN, 'fatbin': enums.CU_JIT_INPUT_FATBINARY, + 'ltoir': enums.CU_JIT_INPUT_LTO_IR, }