Skip to content

Commit

Permalink
Support linking temp ltoir files
Browse files Browse the repository at this point in the history
  • Loading branch information
gmarkall committed Mar 8, 2024
1 parent 4a5586e commit b4df47e
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion pynvjitlink/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,16 @@ def add_file_guess_ext(self, path_or_code):
# To maintain compatibility with the original interface, all strings
# are treated as paths in the filesystem.
if isinstance(path_or_code, str):
return super().add_file_guess_ext(path_or_code)
# 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):
Expand Down Expand Up @@ -187,6 +196,8 @@ def add_data(self, data, kind, name):
return self.add_ptx(data, name)
elif kind == FILE_EXTENSION_MAP["o"]:
fn = self._linker.add_object
elif kind == 'ltoir':
fn = self._linker.add_ltoir
else:
raise LinkerError(f"Don't know how to link {kind}")

Expand Down

0 comments on commit b4df47e

Please sign in to comment.