numba-cuda 0.0.19__tar.gz → 0.0.21__tar.gz

{numba_cuda-0.0.19 → numba_cuda-0.0.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: numba-cuda
-Version: 0.0.19
+Version: 0.0.21
 Summary: CUDA target for Numba
 Author: Anaconda Inc., NVIDIA Corporation
 License: BSD 2-clause
@@ -13,17 +13,21 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numba>=0.59.1
 
+<div align="center"><img src="docs/source/_static/numba-green-icon-rgb.svg" width="200"/></div>
+
 # Numba CUDA Target
 
-An out-of-tree CUDA target for Numba.
+The CUDA target for Numba. Please visit the [official
+documentation](https://nvidia.github.io/numba-cuda) to get started!
+
 
-This contains an entire copy of Numba's CUDA target (the `numba.cuda` module),
-and a mechanism to ensure the code from this module (`numba_cuda.numba.cuda`) is
-used as the `numba.cuda` module instead of the code from the `numba` package.
+To report issues or file feature requests, please use the [issue
+tracker](https://github.com/NVIDIA/numba-cuda/issues).
 
-This is presently in an early state and is published for testing and feedback.
+To raise questions or initiate discussions, please use the [Numba Discourse
+forum](https://numba.discourse.group).
 
-## Building / testing
+## Building from source
 
 Install as an editable install:
 
@@ -31,7 +35,7 @@ Install as an editable install:
 pip install -e .
 ```
 
-Running tests:
+## Running tests
 
 ```
 python -m numba.runtests numba.cuda.tests

numba_cuda-0.0.21/README.md

@@ -0,0 +1,40 @@
+<div align="center"><img src="docs/source/_static/numba-green-icon-rgb.svg" width="200"/></div>
+
+# Numba CUDA Target
+
+The CUDA target for Numba. Please visit the [official
+documentation](https://nvidia.github.io/numba-cuda) to get started!
+
+
+To report issues or file feature requests, please use the [issue
+tracker](https://github.com/NVIDIA/numba-cuda/issues).
+
+To raise questions or initiate discussions, please use the [Numba Discourse
+forum](https://numba.discourse.group).
+
+## Building from source
+
+Install as an editable install:
+
+```
+pip install -e .
+```
+
+## Running tests
+
+```
+python -m numba.runtests numba.cuda.tests
+```
+
+This should discover the`numba.cuda` module from the `numba_cuda` package. You
+can check where `numba.cuda` files are being located by running
+
+```
+python -c "from numba import cuda; print(cuda.__file__)"
+```
+
+which will show a path like:
+
+```
+<path to numba-cuda repo>/numba_cuda/numba/cuda/__init__.py
+```

numba_cuda-0.0.21/numba_cuda/VERSION

@@ -0,0 +1 @@
+0.0.21

{numba_cuda-0.0.19 → numba_cuda-0.0.21}/numba_cuda/numba/cuda/codegen.py

@@ -9,7 +9,6 @@ import os
 import subprocess
 import tempfile
 
-
 CUDA_TRIPLE = 'nvptx64-nvidia-cuda'
 
 
@@ -181,17 +180,7 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
 
         return ltoir
 
-    def get_cubin(self, cc=None):
-        cc = self._ensure_cc(cc)
-
-        cubin = self._cubin_cache.get(cc, None)
-        if cubin:
-            return cubin
-
-        linker = driver.Linker.new(
-            max_registers=self._max_registers, cc=cc, lto=self._lto
-        )
-
+    def _link_all(self, linker, cc, ignore_nonlto=False):
         if linker.lto:
             ltoir = self.get_ltoir(cc=cc)
             linker.add_ltoir(ltoir)
@@ -200,11 +189,44 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
             linker.add_ptx(ptx.encode())
 
         for path in self._linking_files:
-            linker.add_file_guess_ext(path)
+            linker.add_file_guess_ext(path, ignore_nonlto)
         if self.needs_cudadevrt:
-            linker.add_file_guess_ext(get_cudalib('cudadevrt', static=True))
+            linker.add_file_guess_ext(
+                get_cudalib('cudadevrt', static=True), ignore_nonlto
+            )
+
+    def get_cubin(self, cc=None):
+        cc = self._ensure_cc(cc)
 
+        cubin = self._cubin_cache.get(cc, None)
+        if cubin:
+            return cubin
+
+        if self._lto and config.DUMP_ASSEMBLY:
+            linker = driver.Linker.new(
+                max_registers=self._max_registers,
+                cc=cc,
+                additional_flags=["-ptx"],
+                lto=self._lto
+            )
+            # `-ptx` flag is meant to view the optimized PTX for LTO objects.
+            # Non-LTO objects are not passed to linker.
+            self._link_all(linker, cc, ignore_nonlto=True)
+
+            ptx = linker.get_linked_ptx().decode('utf-8')
+
+            print(("ASSEMBLY (AFTER LTO) %s" % self._name).center(80, '-'))
+            print(ptx)
+            print('=' * 80)
+
+        linker = driver.Linker.new(
+            max_registers=self._max_registers,
+            cc=cc,
+            lto=self._lto
+        )
+        self._link_all(linker, cc, ignore_nonlto=False)
         cubin = linker.complete()
+
         self._cubin_cache[cc] = cubin
         self._linkerinfo_cache[cc] = linker.info_log
 

{numba_cuda-0.0.19 → numba_cuda-0.0.21}/numba_cuda/numba/cuda/compiler.py

@@ -1,6 +1,7 @@
 from llvmlite import ir
 from numba.core.typing.templates import ConcreteTemplate
-from numba.core import types, typing, funcdesc, config, compiler, sigutils
+from numba.core import (cgutils, types, typing, funcdesc, config, compiler,
+                        sigutils, utils)
 from numba.core.compiler import (sanitize_compile_result_entries, CompilerBase,
                                  DefaultPassBuilder, Flags, Option,
                                  CompileResult)
@@ -11,7 +12,10 @@ from numba.core.errors import NumbaInvalidConfigWarning
 from numba.core.typed_passes import (IRLegalization, NativeLowering,
                                      AnnotateTypes)
 from warnings import warn
+from numba.cuda import nvvmutils
 from numba.cuda.api import get_current_device
+from numba.cuda.cudadrv import nvvm
+from numba.cuda.descriptor import cuda_target
 from numba.cuda.target import CUDACABICallConv
 
 
@@ -24,6 +28,15 @@ def _nvvm_options_type(x):
         return x
 
 
+def _optional_int_type(x):
+    if x is None:
+        return None
+
+    else:
+        assert isinstance(x, int)
+        return x
+
+
 class CUDAFlags(Flags):
     nvvm_options = Option(
         type=_nvvm_options_type,
@@ -35,6 +48,16 @@ class CUDAFlags(Flags):
         default=None,
         doc="Compute Capability",
     )
+    max_registers = Option(
+        type=_optional_int_type,
+        default=None,
+        doc="Max registers"
+    )
+    lto = Option(
+        type=bool,
+        default=False,
+        doc="Enable Link-time Optimization"
+    )
 
 
 # The CUDACompileResult (CCR) has a specially-defined entry point equal to its
@@ -109,7 +132,11 @@ class CreateLibrary(LoweringPass):
         codegen = state.targetctx.codegen()
         name = state.func_id.func_qualname
         nvvm_options = state.flags.nvvm_options
-        state.library = codegen.create_library(name, nvvm_options=nvvm_options)
+        max_registers = state.flags.max_registers
+        lto = state.flags.lto
+        state.library = codegen.create_library(name, nvvm_options=nvvm_options,
+                                               max_registers=max_registers,
+                                               lto=lto)
         # Enable object caching upfront so that the library can be serialized.
         state.library.enable_object_caching()
 
@@ -152,7 +179,7 @@ class CUDACompiler(CompilerBase):
 @global_compiler_lock
 def compile_cuda(pyfunc, return_type, args, debug=False, lineinfo=False,
                  inline=False, fastmath=False, nvvm_options=None,
-                 cc=None):
+                 cc=None, max_registers=None, lto=False):
     if cc is None:
         raise ValueError('Compute Capability must be supplied')
 
@@ -189,6 +216,8 @@ def compile_cuda(pyfunc, return_type, args, debug=False, lineinfo=False,
     if nvvm_options:
         flags.nvvm_options = nvvm_options
     flags.compute_capability = cc
+    flags.max_registers = max_registers
+    flags.lto = lto
 
     # Run compilation pipeline
     from numba.core.target_extension import target_override
@@ -247,11 +276,155 @@ def cabi_wrap_function(context, lib, fndesc, wrapper_function_name,
         builder, func, restype, argtypes, callargs)
     builder.ret(return_value)
 
+    if config.DUMP_LLVM:
+        utils.dump_llvm(fndesc, wrapper_module)
+
     library.add_ir_module(wrapper_module)
     library.finalize()
     return library
 
 
+def kernel_fixup(kernel, debug):
+    if debug:
+        exc_helper = add_exception_store_helper(kernel)
+
+    # Pass 1 - replace:
+    #
+    #    ret <value>
+    #
+    # with:
+    #
+    #    exc_helper(<value>)
+    #    ret void
+
+    for block in kernel.blocks:
+        for i, inst in enumerate(block.instructions):
+            if isinstance(inst, ir.Ret):
+                old_ret = block.instructions.pop()
+                block.terminator = None
+
+                # The original return's metadata will be set on the new
+                # instructions in order to preserve debug info
+                metadata = old_ret.metadata
+
+                builder = ir.IRBuilder(block)
+                if debug:
+                    status_code = old_ret.operands[0]
+                    exc_helper_call = builder.call(exc_helper, (status_code,))
+                    exc_helper_call.metadata = metadata
+
+                new_ret = builder.ret_void()
+                new_ret.metadata = old_ret.metadata
+
+                # Need to break out so we don't carry on modifying what we are
+                # iterating over. There can only be one return in a block
+                # anyway.
+                break
+
+    # Pass 2: remove stores of null pointer to return value argument pointer
+
+    return_value = kernel.args[0]
+
+    for block in kernel.blocks:
+        remove_list = []
+
+        # Find all stores first
+        for inst in block.instructions:
+            if (isinstance(inst, ir.StoreInstr)
+                    and inst.operands[1] == return_value):
+                remove_list.append(inst)
+
+        # Remove all stores
+        for to_remove in remove_list:
+            block.instructions.remove(to_remove)
+
+    # Replace non-void return type with void return type and remove return
+    # value
+
+    if isinstance(kernel.type, ir.PointerType):
+        new_type = ir.PointerType(ir.FunctionType(ir.VoidType(),
+                                                  kernel.type.pointee.args[1:]))
+    else:
+        new_type = ir.FunctionType(ir.VoidType(), kernel.type.args[1:])
+
+    kernel.type = new_type
+    kernel.return_value = ir.ReturnValue(kernel, ir.VoidType())
+    kernel.args = kernel.args[1:]
+
+    # Mark as a kernel for NVVM
+
+    nvvm.set_cuda_kernel(kernel)
+
+    if config.DUMP_LLVM:
+        print(f"LLVM DUMP: Post kernel fixup {kernel.name}".center(80, '-'))
+        print(kernel.module)
+        print('=' * 80)
+
+
+def add_exception_store_helper(kernel):
+
+    # Create global variables for exception state
+
+    def define_error_gv(postfix):
+        name = kernel.name + postfix
+        gv = cgutils.add_global_variable(kernel.module, ir.IntType(32),
+                                         name)
+        gv.initializer = ir.Constant(gv.type.pointee, None)
+        return gv
+
+    gv_exc = define_error_gv("__errcode__")
+    gv_tid = []
+    gv_ctaid = []
+    for i in 'xyz':
+        gv_tid.append(define_error_gv("__tid%s__" % i))
+        gv_ctaid.append(define_error_gv("__ctaid%s__" % i))
+
+    # Create exception store helper function
+
+    helper_name = kernel.name + "__exc_helper__"
+    helper_type = ir.FunctionType(ir.VoidType(), (ir.IntType(32),))
+    helper_func = ir.Function(kernel.module, helper_type, helper_name)
+
+    block = helper_func.append_basic_block(name="entry")
+    builder = ir.IRBuilder(block)
+
+    # Implement status check / exception store logic
+
+    status_code = helper_func.args[0]
+    call_conv = cuda_target.target_context.call_conv
+    status = call_conv._get_return_status(builder, status_code)
+
+    # Check error status
+    with cgutils.if_likely(builder, status.is_ok):
+        builder.ret_void()
+
+    with builder.if_then(builder.not_(status.is_python_exc)):
+        # User exception raised
+        old = ir.Constant(gv_exc.type.pointee, None)
+
+        # Use atomic cmpxchg to prevent rewriting the error status
+        # Only the first error is recorded
+
+        xchg = builder.cmpxchg(gv_exc, old, status.code,
+                               'monotonic', 'monotonic')
+        changed = builder.extract_value(xchg, 1)
+
+        # If the xchange is successful, save the thread ID.
+        sreg = nvvmutils.SRegBuilder(builder)
+        with builder.if_then(changed):
+            for dim, ptr, in zip("xyz", gv_tid):
+                val = sreg.tid(dim)
+                builder.store(val, ptr)
+
+            for dim, ptr, in zip("xyz", gv_ctaid):
+                val = sreg.ctaid(dim)
+                builder.store(val, ptr)
+
+    builder.ret_void()
+
+    return helper_func
+
+
 @global_compiler_lock
 def compile(pyfunc, sig, debug=None, lineinfo=False, device=True,
             fastmath=False, cc=None, opt=None, abi="c", abi_info=None,
@@ -347,13 +520,10 @@ def compile(pyfunc, sig, debug=None, lineinfo=False, device=True,
             lib = cabi_wrap_function(tgt, lib, cres.fndesc, wrapper_name,
                                      nvvm_options)
     else:
-        code = pyfunc.__code__
-        filename = code.co_filename
-        linenum = code.co_firstlineno
-
-        lib, kernel = tgt.prepare_cuda_kernel(cres.library, cres.fndesc, debug,
-                                              lineinfo, nvvm_options, filename,
-                                              linenum)
+        lib = cres.library
+        kernel = lib.get_function(cres.fndesc.llvm_func_name)
+        lib._entry_name = cres.fndesc.llvm_func_name
+        kernel_fixup(kernel, debug)
 
     if lto:
         code = lib.get_ltoir(cc=cc)

{numba_cuda-0.0.19 → numba_cuda-0.0.21}/numba_cuda/numba/cuda/cuda_paths.py

@@ -310,7 +310,9 @@ def get_conda_include_dir():
         # though usually it shouldn't.
         include_dir = os.path.join(sys.prefix, 'include')
 
-    if os.path.exists(include_dir):
+    if (os.path.exists(include_dir) and os.path.isdir(include_dir)
+            and os.path.exists(os.path.join(include_dir,
+                                            'cuda_device_runtime_api.h'))):
         return include_dir
     return
 

{numba_cuda-0.0.19 → numba_cuda-0.0.21}/numba_cuda/numba/cuda/cudadrv/driver.py

@@ -21,6 +21,9 @@ import threading
 import traceback
 import asyncio
 import pathlib
+import subprocess
+import tempfile
+import re
 from itertools import product
 from abc import ABCMeta, abstractmethod
 from ctypes import (c_int, byref, c_size_t, c_char, c_char_p, addressof,
@@ -36,7 +39,7 @@ from .error import CudaSupportError, CudaDriverError
 from .drvapi import API_PROTOTYPES
 from .drvapi import cu_occupancy_b2d_size, cu_stream_callback_pyobj, cu_uuid
 from .mappings import FILE_EXTENSION_MAP
-from .linkable_code import LinkableCode
+from .linkable_code import LinkableCode, LTOIR, Fatbin, Object
 from numba.cuda.cudadrv import enums, drvapi, nvrtc
 
 USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING
@@ -2683,12 +2686,18 @@ class Linker(metaclass=ABCMeta):
             cu = f.read()
         self.add_cu(cu, os.path.basename(path))
 
-    def add_file_guess_ext(self, path_or_code):
+    def add_file_guess_ext(self, path_or_code, ignore_nonlto=False):
         """
         Add a file or LinkableCode object to the link. If a file is
         passed, the type will be inferred from the extension. A LinkableCode
         object represents a file already in memory.
+
+        When `ignore_nonlto` is set to true, do not add code that will not
+        be LTO-ed in the linking process. This is useful in inspecting the
+        LTO-ed portion of the PTX when linker is added with objects that can be
+        both LTO-ed and not LTO-ed.
         """
+
         if isinstance(path_or_code, str):
             ext = pathlib.Path(path_or_code).suffix
             if ext == '':
@@ -2704,6 +2713,26 @@ class Linker(metaclass=ABCMeta):
                         "Don't know how to link file with extension "
                         f"{ext}"
                     )
+
+                if ignore_nonlto:
+                    warn_and_return = False
+                    if kind in (
+                        FILE_EXTENSION_MAP["fatbin"], FILE_EXTENSION_MAP["o"]
+                    ):
+                        entry_types = inspect_obj_content(path_or_code)
+                        if "nvvm" not in entry_types:
+                            warn_and_return = True
+                    elif kind != FILE_EXTENSION_MAP["ltoir"]:
+                        warn_and_return = True
+
+                    if warn_and_return:
+                        warnings.warn(
+                            f"Not adding {path_or_code} as it is not "
+                            "optimizable at link time, and `ignore_nonlto == "
+                            "True`."
+                        )
+                        return
+
                 self.add_file(path_or_code, kind)
             return
         else:
@@ -2716,6 +2745,25 @@ class Linker(metaclass=ABCMeta):
             if path_or_code.kind == "cu":
                 self.add_cu(path_or_code.data, path_or_code.name)
             else:
+                if ignore_nonlto:
+                    warn_and_return = False
+                    if isinstance(path_or_code, (Fatbin, Object)):
+                        with tempfile.NamedTemporaryFile("w") as fp:
+                            fp.write(path_or_code.data)
+                            entry_types = inspect_obj_content(fp.name)
+                        if "nvvm" not in entry_types:
+                            warn_and_return = True
+                    elif not isinstance(path_or_code, LTOIR):
+                        warn_and_return = True
+
+                    if warn_and_return:
+                        warnings.warn(
+                            f"Not adding {path_or_code.name} as it is not "
+                            "optimizable at link time, and `ignore_nonlto == "
+                            "True`."
+                        )
+                        return
+
                 self.add_data(
                     path_or_code.data, path_or_code.kind, path_or_code.name
                 )
@@ -3065,6 +3113,28 @@ class PyNvJitLinker(Linker):
         name = pathlib.Path(path).name
         self.add_data(data, kind, 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
+
+        program, log = nvrtc.compile(cu, name, cc, ltoir=self.lto)
+
+        if not self.lto and config.DUMP_ASSEMBLY:
+            print(("ASSEMBLY %s" % name).center(80, "-"))
+            print(program)
+            print("=" * 80)
+
+        suffix = ".ltoir" if self.lto else ".ptx"
+        program_name = os.path.splitext(name)[0] + suffix
+        # Link the program's PTX or LTOIR using the normal linker mechanism
+        if self.lto:
+            self.add_ltoir(program, program_name)
+        else:
+            self.add_ptx(program.encode(), program_name)
+
     def add_data(self, data, kind, name):
         if kind == FILE_EXTENSION_MAP["cubin"]:
             fn = self._linker.add_cubin
@@ -3086,6 +3156,12 @@ class PyNvJitLinker(Linker):
         except NvJitLinkError as e:
             raise LinkerError from e
 
+    def get_linked_ptx(self):
+        try:
+            return self._linker.get_linked_ptx()
+        except NvJitLinkError as e:
+            raise LinkerError from e
+
     def complete(self):
         try:
             return self._linker.get_linked_cubin()
@@ -3361,3 +3437,28 @@ def get_version():
     Return the driver version as a tuple of (major, minor)
     """
     return driver.get_version()
+
+
+def inspect_obj_content(objpath: str):
+    """
+    Given path to a fatbin or object, use `cuobjdump` to examine its content
+    Return the set of entries in the object.
+    """
+    code_types :set[str] = set()
+
+    try:
+        out = subprocess.run(["cuobjdump", objpath], check=True,
+                             capture_output=True)
+    except FileNotFoundError as e:
+        msg = ("cuobjdump has not been found. You may need "
+               "to install the CUDA toolkit and ensure that "
+               "it is available on your PATH.\n")
+        raise RuntimeError(msg) from e
+
+    objtable = out.stdout.decode('utf-8')
+    entry_pattern = r"Fatbin (.*) code"
+    for line in objtable.split("\n"):
+        if match := re.match(entry_pattern, line):
+            code_types.add(match.group(1))
+
+    return code_types

{numba_cuda-0.0.19 → numba_cuda-0.0.21}/numba_cuda/numba/cuda/cudadrv/enums.py

@@ -55,7 +55,7 @@ CUDA_ERROR_INVALID_HANDLE = 400
 CUDA_ERROR_ILLEGAL_STATE = 401
 CUDA_ERROR_NOT_FOUND = 500
 CUDA_ERROR_NOT_READY = 600
-CUDA_ERROR_LAUNCH_FAILED = 700
+CUDA_ERROR_ILLEGAL_ADDRESS = 700
 CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES = 701
 CUDA_ERROR_LAUNCH_TIMEOUT = 702
 CUDA_ERROR_LAUNCH_INCOMPATIBLE_TEXTURING = 703

{numba_cuda-0.0.19 → numba_cuda-0.0.21}/numba_cuda/numba/cuda/cudadrv/nvrtc.py

@@ -61,6 +61,14 @@ class NVRTC:
     NVVM interface. Initialization is protected by a lock and uses the standard
     (for Numba) open_cudalib function to load the NVRTC library.
     """
+
+    _CU12ONLY_PROTOTYPES = {
+        # nvrtcResult nvrtcGetLTOIRSize(nvrtcProgram prog, size_t *ltoSizeRet);
+        "nvrtcGetLTOIRSize": (nvrtc_result, nvrtc_program, POINTER(c_size_t)),
+        # nvrtcResult nvrtcGetLTOIR(nvrtcProgram prog, char *lto);
+        "nvrtcGetLTOIR": (nvrtc_result, nvrtc_program, c_char_p)
+    }
+
     _PROTOTYPES = {
         # nvrtcResult nvrtcVersion(int *major, int *minor)
         'nvrtcVersion': (nvrtc_result, POINTER(c_int), POINTER(c_int)),
@@ -110,6 +118,10 @@ class NVRTC:
                     cls.__INSTANCE = None
                     raise NvrtcSupportError("NVRTC cannot be loaded") from e
 
+                from numba.cuda.cudadrv.runtime import get_version
+                if get_version() >= (12, 0):
+                    inst._PROTOTYPES |= inst._CU12ONLY_PROTOTYPES
+
                 # Find & populate functions
                 for name, proto in inst._PROTOTYPES.items():
                     func = getattr(lib, name)
@@ -208,10 +220,22 @@ class NVRTC:
 
         return ptx.value.decode()
 
+    def get_lto(self, program):
+        """
+        Get the compiled LTOIR as a Python bytes object.
+        """
+        lto_size = c_size_t()
+        self.nvrtcGetLTOIRSize(program.handle, byref(lto_size))
+
+        lto = b" " * lto_size.value
+        self.nvrtcGetLTOIR(program.handle, lto)
+
+        return lto
 
-def compile(src, name, cc):
+
+def compile(src, name, cc, ltoir=False):
     """
-    Compile a CUDA C/C++ source to PTX for a given compute capability.
+    Compile a CUDA C/C++ source to PTX or LTOIR for a given compute capability.
 
     :param src: The source code to compile
     :type src: str
@@ -219,6 +243,8 @@ def compile(src, name, cc):
     :type name: str
     :param cc: A tuple ``(major, minor)`` of the compute capability
     :type cc: tuple
+    :param ltoir: Compile into LTOIR if True, otherwise into PTX
+    :type ltoir: bool
     :return: The compiled PTX and compilation log
     :rtype: tuple
     """
@@ -242,6 +268,9 @@ def compile(src, name, cc):
     numba_include = f'-I{numba_cuda_path}'
     options = [arch, *cuda_include, numba_include, '-rdc', 'true']
 
+    if ltoir:
+        options.append("-dlto")
+
     if nvrtc.get_version() < (12, 0):
         options += ["-std=c++17"]
 
@@ -261,5 +290,9 @@ def compile(src, name, cc):
         msg = (f"NVRTC log messages whilst compiling {name}:\n\n{log}")
         warnings.warn(msg)
 
-    ptx = nvrtc.get_ptx(program)
-    return ptx, log
+    if ltoir:
+        ltoir = nvrtc.get_lto(program)
+        return ltoir, log
+    else:
+        ptx = nvrtc.get_ptx(program)
+        return ptx, log