Add nvjitlink get_linked_ltoir (#2337)

Author: ZzEeKkAa

cuda_bindings/docs/source/module/nvjitlink.rst

@@ -25,6 +25,8 @@ NvJitLink defines the following functions for linking code objects and querying
 .. autofunction:: cuda.bindings.nvjitlink.get_linked_cubin
 .. autofunction:: cuda.bindings.nvjitlink.get_linked_ptx_size
 .. autofunction:: cuda.bindings.nvjitlink.get_linked_ptx
+.. autofunction:: cuda.bindings.nvjitlink.get_linked_ltoir_size
+.. autofunction:: cuda.bindings.nvjitlink.get_linked_ltoir
 .. autofunction:: cuda.bindings.nvjitlink.get_error_log_size
 .. autofunction:: cuda.bindings.nvjitlink.get_error_log
 .. autofunction:: cuda.bindings.nvjitlink.get_info_log_size
@@ -65,6 +67,33 @@ Types
     .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_FINALIZE
 
 
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_NULL_INPUT
+
+
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_INCOMPATIBLE_OPTIONS
+
+
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_INCORRECT_INPUT_TYPE
+
+
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_ARCH_MISMATCH
+
+
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_OUTDATED_LIBRARY
+
+
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_MISSING_FATBIN
+
+
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_UNRECOGNIZED_ARCH
+
+
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_UNSUPPORTED_ARCH
+
+
+    .. autoattribute:: cuda.bindings.nvjitlink.Result.ERROR_LTO_NOT_ENABLED
+
+
 .. autoclass:: cuda.bindings.nvjitlink.InputType
 
     .. autoattribute:: cuda.bindings.nvjitlink.InputType.NONE

cuda_bindings/tests/test_nvjitlink.py

@@ -70,6 +70,12 @@ def check_nvjitlink_usable():
     return inner_nvjitlink._inspect_function_pointer("__nvJitLinkVersion") != 0
 
 
+def check_nvjitlink_get_linked_ltoir_usable():
+    from cuda.bindings._internal import nvjitlink as inner_nvjitlink
+
+    return inner_nvjitlink._inspect_function_pointer("__nvJitLinkGetLinkedLTOIR") != 0
+
+
 pytestmark = pytest.mark.skipif(
     not check_nvjitlink_usable(), reason="nvJitLink not usable, maybe not installed or too old (<12.3)"
 )
@@ -183,6 +189,21 @@ def test_get_linked_ptx(arch, get_dummy_ltoir):
         assert len(ptx) == ptx_size
 
 
+@pytest.mark.parametrize("arch", ARCHITECTURES)
+@pytest.mark.skipif(
+    not check_nvjitlink_get_linked_ltoir_usable(),
+    reason="nvJitLinkGetLinkedLTOIR not available in installed nvJitLink",
+)
+def test_get_linked_ltoir(arch, get_dummy_ltoir):
+    with nvjitlink_session(2, [f"-arch={arch}", "-lto"]) as handle:
+        nvjitlink.add_data(handle, nvjitlink.InputType.LTOIR, get_dummy_ltoir, len(get_dummy_ltoir), "test_data")
+        nvjitlink.complete(handle)
+        ltoir_size = nvjitlink.get_linked_ltoir_size(handle)
+        ltoir = bytearray(ltoir_size)
+        nvjitlink.get_linked_ltoir(handle, ltoir)
+        assert len(ltoir) == ltoir_size
+
+
 def test_package_version():
     ver = nvjitlink.version()
     assert len(ver) == 2