numba-cuda 0.10.1__tar.gz → 0.12.1__tar.gz

{numba_cuda-0.10.1 → numba_cuda-0.12.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: numba-cuda
-Version: 0.10.1
+Version: 0.12.1
 Summary: CUDA target for Numba
 Author: Anaconda Inc., NVIDIA Corporation
 License: BSD 2-clause

numba_cuda-0.12.1/numba_cuda/VERSION

@@ -0,0 +1 @@
+0.12.1

{numba_cuda-0.10.1/numba_cuda/numba/cuda → numba_cuda-0.12.1/numba_cuda/numba/cuda/_internal}/cuda_bf16.py

@@ -2,7 +2,7 @@
 # Generator Information:
 # Ast_canopy version: 0.3.0
 # Numbast version: 0.3.0
-# Generation command: /home/wangm/numbast/numbast/src/numbast/__main__.py --cfg-path configs/cuda_bf16.yml --output-dir numba_cuda/numba/cuda/
+# Generation command: /home/wangm/numbast/numbast/src/numbast/__main__.py --cfg-path configs/cuda_bf16.yml --output-dir numba_cuda/numba/cuda/_internal
 # Static binding generator parameters: {'cfg_path': 'configs/cuda_bf16.yml', 'output_dir': 'numba_cuda/numba/cuda/', 'entry_point': None, 'retain': None, 'types': None, 'datamodels': None, 'compute_capability': None, 'run_ruff_format': True}
 # Config file path (relative to the path of the generated binding): ../../../../configs/cuda_bf16.yml
 # Cudatoolkit version: (12, 8)

{numba_cuda-0.10.1 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/api.py

@@ -10,6 +10,7 @@ import numpy as np
 from .cudadrv import devicearray, devices, driver
 from numba.core import config
 from numba.cuda.api_util import prepare_shape_strides_dtype
+from numba.cuda.cudadrv.runtime import get_version
 
 # NDarray device helper
 
@@ -95,6 +96,18 @@ def is_float16_supported():
     return True
 
 
+def is_bfloat16_supported():
+    """Whether bfloat16 are supported.
+
+    bfloat16 are only supported on devices with compute capability >= 8.0 and cuda version >= 12.0
+    """
+    cuda_version = get_version()
+    return current_context().device.supports_bfloat16 and cuda_version >= (
+        12,
+        0,
+    )
+
+
 @require_context
 def to_device(obj, stream=0, copy=True, to=None):
     """to_device(obj, stream=0, copy=True, to=None)

numba_cuda-0.12.1/numba_cuda/numba/cuda/bf16.py

@@ -0,0 +1,112 @@
+from numba.cuda._internal.cuda_bf16 import (
+    _type_class___nv_bfloat16,
+    nv_bfloat16 as bfloat16,
+    htrunc,
+    hceil,
+    hfloor,
+    hrint,
+    hsqrt,
+    hrsqrt,
+    hrcp,
+    hlog,
+    hlog2,
+    hlog10,
+    hcos,
+    hsin,
+    hexp,
+    hexp2,
+    hexp10,
+    htanh,
+    htanh_approx,
+)
+from numba.extending import overload
+
+import math
+
+
+def _make_unary(a, func):
+    if isinstance(a, _type_class___nv_bfloat16):
+        return lambda a: func(a)
+
+
+# Bind low++ bindings to math APIs
+@overload(math.trunc, target="cuda")
+def trunc_ol(a):
+    return _make_unary(a, htrunc)
+
+
+@overload(math.ceil, target="cuda")
+def ceil_ol(a):
+    return _make_unary(a, hceil)
+
+
+@overload(math.floor, target="cuda")
+def floor_ol(a):
+    return _make_unary(a, hfloor)
+
+
+@overload(math.sqrt, target="cuda")
+def sqrt_ol(a):
+    return _make_unary(a, hsqrt)
+
+
+@overload(math.log, target="cuda")
+def log_ol(a):
+    return _make_unary(a, hlog)
+
+
+@overload(math.log10, target="cuda")
+def log10_ol(a):
+    return _make_unary(a, hlog10)
+
+
+@overload(math.cos, target="cuda")
+def cos_ol(a):
+    return _make_unary(a, hcos)
+
+
+@overload(math.sin, target="cuda")
+def sin_ol(a):
+    return _make_unary(a, hsin)
+
+
+@overload(math.tanh, target="cuda")
+def tanh_ol(a):
+    return _make_unary(a, htanh)
+
+
+@overload(math.exp, target="cuda")
+def exp_ol(a):
+    return _make_unary(a, hexp)
+
+
+try:
+    from math import exp2
+
+    @overload(exp2, target="cuda")
+    def exp2_ol(a):
+        return _make_unary(a, hexp2)
+except ImportError:
+    pass
+
+
+__all__ = [
+    "bfloat16",
+    "htrunc",
+    "hceil",
+    "hfloor",
+    "hrint",
+    "hsqrt",
+    "hrsqrt",
+    "hrcp",
+    "hlog",
+    "hlog2",
+    "hlog10",
+    "hcos",
+    "hsin",
+    "htanh",
+    "htanh_approx",
+    "hexp",
+    "hexp2",
+    "hexp10",
+]

{numba_cuda-0.10.1 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/cg.py

@@ -23,6 +23,7 @@ def _this_grid(typingctx):
     sig = signature(grid_group)
 
     def codegen(context, builder, sig, args):
+        context.active_code_library.use_cooperative = True
         one = context.get_constant(types.int32, 1)
         mod = builder.module
         return builder.call(
@@ -45,6 +46,7 @@ def _grid_group_sync(typingctx, group):
     sig = signature(types.int32, group)
 
     def codegen(context, builder, sig, args):
+        context.active_code_library.use_cooperative = True
         flags = context.get_constant(types.int32, 0)
         mod = builder.module
         return builder.call(

{numba_cuda-0.10.1 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/codegen.py

@@ -5,6 +5,7 @@ from numba.core.codegen import Codegen, CodeLibrary
 from .cudadrv import devices, driver, nvvm, runtime
 from numba.cuda.cudadrv.libs import get_cudalib
 from numba.cuda.cudadrv.linkable_code import LinkableCode
+from numba.cuda.runtime.nrt import NRT_LIBRARY
 
 import os
 import subprocess
@@ -57,6 +58,59 @@ def disassemble_cubin_for_cfg(cubin):
     return run_nvdisasm(cubin, flags)
 
 
+class ExternalCodeLibrary(CodeLibrary):
+    """Holds code produced externally, for linking with generated code."""
+
+    def __init__(self, codegen, name):
+        super().__init__(codegen, name)
+        # Files to link
+        self._linking_files = set()
+        # Setup and teardown functions for the module.
+        # The order is determined by the order they are added to the codelib.
+        self._setup_functions = []
+        self._teardown_functions = []
+
+        self.use_cooperative = False
+
+    @property
+    def modules(self):
+        # There are no LLVM IR modules in an ExternalCodeLibrary
+        return set()
+
+    def add_linking_file(self, path_or_obj):
+        # Adding new files after finalization is prohibited, in case the list
+        # of libraries has already been added to another code library; the
+        # newly-added files would be omitted from their linking process.
+        self._raise_if_finalized()
+
+        if isinstance(path_or_obj, LinkableCode):
+            if path_or_obj.setup_callback:
+                self._setup_functions.append(path_or_obj.setup_callback)
+            if path_or_obj.teardown_callback:
+                self._teardown_functions.append(path_or_obj.teardown_callback)
+
+        self._linking_files.add(path_or_obj)
+
+    def add_ir_module(self, module):
+        raise NotImplementedError("Cannot add LLVM IR to external code")
+
+    def add_linking_library(self, library):
+        raise NotImplementedError("Cannot add libraries to external code")
+
+    def finalize(self):
+        self._raise_if_finalized()
+        self._finalized = True
+
+    def get_asm_str(self):
+        raise NotImplementedError("No assembly for external code")
+
+    def get_llvm_str(self):
+        raise NotImplementedError("No LLVM IR for external code")
+
+    def get_function(self, name):
+        raise NotImplementedError("Cannot get function from external code")
+
+
 class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
     """
     The CUDACodeLibrary generates PTX, SASS, cubins for multiple different
@@ -129,6 +183,8 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         self._nvvm_options = nvvm_options
         self._entry_name = entry_name
 
+        self.use_cooperative = False
+
     @property
     def llvm_strs(self):
         if self._llvm_strs is None:
@@ -297,6 +353,10 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         self._raise_if_finalized()
 
         self._linking_libraries.add(library)
+        self._linking_files.update(library._linking_files)
+        self._setup_functions.extend(library._setup_functions)
+        self._teardown_functions.extend(library._teardown_functions)
+        self.use_cooperative |= library.use_cooperative
 
     def add_linking_file(self, path_or_obj):
         if isinstance(path_or_obj, LinkableCode):
@@ -362,9 +422,17 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         but loaded functions are discarded. They are recreated when needed
         after deserialization.
         """
+        nrt = False
         if self._linking_files:
-            msg = "Cannot pickle CUDACodeLibrary with linking files"
-            raise RuntimeError(msg)
+            if (
+                len(self._linking_files) == 1
+                and NRT_LIBRARY in self._linking_files
+            ):
+                nrt = True
+            else:
+                msg = "Cannot pickle CUDACodeLibrary with linking files"
+                raise RuntimeError(msg)
+
         if not self._finalized:
             raise RuntimeError("Cannot pickle unfinalized CUDACodeLibrary")
         return dict(
@@ -378,6 +446,8 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
             max_registers=self._max_registers,
             nvvm_options=self._nvvm_options,
             needs_cudadevrt=self.needs_cudadevrt,
+            nrt=nrt,
+            use_cooperative=self.use_cooperative,
         )
 
     @classmethod
@@ -393,6 +463,8 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         max_registers,
         nvvm_options,
         needs_cudadevrt,
+        nrt,
+        use_cooperative,
     ):
         """
         Rebuild an instance.
@@ -407,8 +479,11 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         instance._max_registers = max_registers
         instance._nvvm_options = nvvm_options
         instance.needs_cudadevrt = needs_cudadevrt
+        instance.use_cooperative = use_cooperative
 
         instance._finalized = True
+        if nrt:
+            instance._linking_files = {NRT_LIBRARY}
 
         return instance
 

{numba_cuda-0.10.1 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/compiler.py

@@ -1,5 +1,4 @@
 from llvmlite import ir
-from numba.core.typing.templates import ConcreteTemplate
 from numba.core import ir as numba_ir
 from numba.core import (
     cgutils,
@@ -37,6 +36,7 @@ from numba.core.typed_passes import (
 from warnings import warn
 from numba.cuda import nvvmutils
 from numba.cuda.api import get_current_device
+from numba.cuda.codegen import ExternalCodeLibrary
 from numba.cuda.cudadrv import nvvm
 from numba.cuda.descriptor import cuda_target
 from numba.cuda.target import CUDACABICallConv
@@ -797,33 +797,39 @@ def compile_ptx_for_current_device(
     )
 
 
-def declare_device_function(name, restype, argtypes, link):
-    return declare_device_function_template(name, restype, argtypes, link).key
-
-
-def declare_device_function_template(name, restype, argtypes, link):
+def declare_device_function(name, restype, argtypes, link, use_cooperative):
     from .descriptor import cuda_target
 
     typingctx = cuda_target.typing_context
     targetctx = cuda_target.target_context
     sig = typing.signature(restype, *argtypes)
-    extfn = ExternFunction(name, sig, link)
 
-    class device_function_template(ConcreteTemplate):
-        key = extfn
-        cases = [sig]
+    # extfn is the descriptor used to call the function from Python code, and
+    # is used as the key for typing and lowering.
+    extfn = ExternFunction(name, sig)
 
-    fndesc = funcdesc.ExternalFunctionDescriptor(
-        name=name, restype=restype, argtypes=argtypes
-    )
+    # Typing
+    device_function_template = typing.make_concrete_template(name, extfn, [sig])
     typingctx.insert_user_function(extfn, device_function_template)
-    targetctx.insert_user_function(extfn, fndesc)
+
+    # Lowering
+    lib = ExternalCodeLibrary(f"{name}_externals", targetctx.codegen())
+    for file in link:
+        lib.add_linking_file(file)
+    lib.use_cooperative = use_cooperative
+
+    # ExternalFunctionDescriptor provides a lowering implementation for calling
+    # external functions
+    fndesc = funcdesc.ExternalFunctionDescriptor(name, restype, argtypes)
+    targetctx.insert_user_function(extfn, fndesc, libs=(lib,))
 
     return device_function_template
 
 
 class ExternFunction:
-    def __init__(self, name, sig, link):
+    """A descriptor that can be used to call the external function from within
+    a Python kernel."""
+
+    def __init__(self, name, sig):
         self.name = name
         self.sig = sig
-        self.link = link

{numba_cuda-0.10.1 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/cudadecl.py

@@ -1,5 +1,5 @@
 import operator
-from numba.core import types
+from numba.core import errors, types
 from numba.core.typing.npydecl import (
     parse_dtype,
     parse_shape,
@@ -21,7 +21,7 @@ from numba.core.typing.templates import (
 from numba.cuda.types import dim3
 from numba.core.typeconv import Conversion
 from numba import cuda
-from numba.cuda.compiler import declare_device_function_template
+from numba.cuda.compiler import declare_device_function
 
 registry = Registry()
 register = registry.register
@@ -33,7 +33,7 @@ register_number_classes(register_global)
 
 class Cuda_array_decl(CallableTemplate):
     def generic(self):
-        def typer(shape, dtype):
+        def typer(shape, dtype, alignment=None):
             # Only integer literals and tuples of integer literals are valid
             # shapes
             if isinstance(shape, types.Integer):
@@ -47,6 +47,16 @@ class Cuda_array_decl(CallableTemplate):
             else:
                 return None
 
+            if alignment is not None:
+                permitted = (types.IntegerLiteral, types.NoneType)
+                if not isinstance(alignment, permitted):
+                    msg = "alignment must be a constant integer"
+                    raise errors.RequireLiteralValue(msg)
+
+            # N.B. We don't use alignment for typing; it's not part of
+            #      types.Array.  The value supplied to the array declaration
+            #      is handled in the lowering.
+
             ndim = parse_shape(shape)
             nb_dtype = parse_dtype(dtype)
             if nb_dtype is not None and ndim is not None:
@@ -412,15 +422,19 @@ _genfp16_binary_operator(operator.itruediv)
 
 def _resolve_wrapped_unary(fname):
     link = tuple()
-    decl = declare_device_function_template(
-        f"__numba_wrapper_{fname}", types.float16, (types.float16,), link
+    decl = declare_device_function(
+        f"__numba_wrapper_{fname}",
+        types.float16,
+        (types.float16,),
+        link,
+        use_cooperative=False,
     )
     return types.Function(decl)
 
 
 def _resolve_wrapped_binary(fname):
     link = tuple()
-    decl = declare_device_function_template(
+    decl = declare_device_function(
         f"__numba_wrapper_{fname}",
         types.float16,
         (
@@ -428,6 +442,7 @@ def _resolve_wrapped_binary(fname):
             types.float16,
         ),
         link,
+        use_cooperative=False,
     )
     return types.Function(decl)
 

{numba_cuda-0.10.1 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/cudadrv/driver.py

@@ -49,7 +49,7 @@ 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, LTOIR, Fatbin, Object
-from numba.cuda.utils import _readenv
+from numba.cuda.utils import _readenv, cached_file_read
 from numba.cuda.cudadrv import enums, drvapi, nvrtc
 
 try:
@@ -714,6 +714,10 @@ class Device(object):
     def supports_float16(self):
         return self.compute_capability >= (5, 3)
 
+    @property
+    def supports_bfloat16(self):
+        return self.compute_capability >= (8, 0)
+
 
 def met_requirement_for_device(device):
     if device.compute_capability < MIN_REQUIRED_CC:
@@ -2797,13 +2801,16 @@ class Linker(metaclass=ABCMeta):
         ptx_name = os.path.splitext(name)[0] + ".ptx"
         self.add_ptx(ptx.encode(), ptx_name)
 
+    @abstractmethod
+    def add_data(self, data, kind, name):
+        """Add in-memory data to the link"""
+
     @abstractmethod
     def add_file(self, path, kind):
         """Add code from a file to the link"""
 
     def add_cu_file(self, path):
-        with open(path, "rb") as f:
-            cu = f.read()
+        cu = cached_file_read(path, how="rb")
         self.add_cu(cu, os.path.basename(path))
 
     def add_file_guess_ext(self, path_or_code, ignore_nonlto=False):
@@ -2948,6 +2955,10 @@ class MVCLinker(Linker):
         except CubinLinkerError as e:
             raise LinkerError from e
 
+    def add_data(self, data, kind, name):
+        msg = "Adding in-memory data unsupported in the MVC linker"
+        raise LinkerError(msg)
+
     def add_file(self, path, kind):
         try:
             from cubinlinker import CubinLinkerError
@@ -2955,8 +2966,7 @@ class MVCLinker(Linker):
             raise ImportError(_MVC_ERROR_MESSAGE) from err
 
         try:
-            with open(path, "rb") as f:
-                data = f.read()
+            data = cached_file_read(path, how="rb")
         except FileNotFoundError:
             raise LinkerError(f"{path} not found")
 
@@ -3046,17 +3056,32 @@ class CtypesLinker(Linker):
     def error_log(self):
         return self.linker_errors_buf.value.decode("utf8")
 
-    def add_ptx(self, ptx, name="<cudapy-ptx>"):
-        ptxbuf = c_char_p(ptx)
-        namebuf = c_char_p(name.encode("utf8"))
-        self._keep_alive += [ptxbuf, namebuf]
+    def add_cubin(self, cubin, name="<unnamed-cubin>"):
+        return self._add_data(enums.CU_JIT_INPUT_CUBIN, cubin, name)
+
+    def add_ptx(self, ptx, name="<unnamed-ptx>"):
+        return self._add_data(enums.CU_JIT_INPUT_PTX, ptx, name)
+
+    def add_object(self, object_, name="<unnamed-object>"):
+        return self._add_data(enums.CU_JIT_INPUT_OBJECT, object_, name)
+
+    def add_fatbin(self, fatbin, name="<unnamed-fatbin>"):
+        return self._add_data(enums.CU_JIT_INPUT_FATBINARY, fatbin, name)
+
+    def add_library(self, library, name="<unnamed-library>"):
+        return self._add_data(enums.CU_JIT_INPUT_LIBRARY, library, name)
+
+    def _add_data(self, input_type, data, name):
+        data_buffer = c_char_p(data)
+        name_buffer = c_char_p(name.encode("utf8"))
+        self._keep_alive += [data_buffer, name_buffer]
         try:
             driver.cuLinkAddData(
                 self.handle,
-                enums.CU_JIT_INPUT_PTX,
-                ptxbuf,
-                len(ptx),
-                namebuf,
+                input_type,
+                data_buffer,
+                len(data),
+                name_buffer,
                 0,
                 None,
                 None,
@@ -3064,6 +3089,28 @@ class CtypesLinker(Linker):
         except CudaAPIError as e:
             raise LinkerError("%s\n%s" % (e, self.error_log))
 
+    def add_data(self, data, kind, name=None):
+        # We pass the name as **kwargs to ensure the default name for the input
+        # type is used if none is supplied
+        kws = {}
+        if name is not None:
+            kws["name"] = name
+
+        if kind == FILE_EXTENSION_MAP["cubin"]:
+            self.add_cubin(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["fatbin"]:
+            self.add_fatbin(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["a"]:
+            self.add_library(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["ptx"]:
+            self.add_ptx(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["o"]:
+            self.add_object(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["ltoir"]:
+            raise LinkerError("Ctypes linker cannot link LTO-IR")
+        else:
+            raise LinkerError(f"Don't know how to link {kind}")
+
     def add_file(self, path, kind):
         pathbuf = c_char_p(path.encode("utf8"))
         self._keep_alive.append(pathbuf)
@@ -3151,17 +3198,58 @@ class CudaPythonLinker(Linker):
     def error_log(self):
         return self.linker_errors_buf.decode("utf8")
 
-    def add_ptx(self, ptx, name="<cudapy-ptx>"):
-        namebuf = name.encode("utf8")
-        self._keep_alive += [ptx, namebuf]
+    def add_cubin(self, cubin, name="<unnamed-cubin>"):
+        input_type = binding.CUjitInputType.CU_JIT_INPUT_CUBIN
+        return self._add_data(input_type, cubin, name)
+
+    def add_ptx(self, ptx, name="<unnamed-ptx>"):
+        input_type = binding.CUjitInputType.CU_JIT_INPUT_PTX
+        return self._add_data(input_type, ptx, name)
+
+    def add_object(self, object_, name="<unnamed-object>"):
+        input_type = binding.CUjitInputType.CU_JIT_INPUT_OBJECT
+        return self._add_data(input_type, object_, name)
+
+    def add_fatbin(self, fatbin, name="<unnamed-fatbin>"):
+        input_type = binding.CUjitInputType.CU_JIT_INPUT_FATBINARY
+        return self._add_data(input_type, fatbin, name)
+
+    def add_library(self, library, name="<unnamed-library>"):
+        input_type = binding.CUjitInputType.CU_JIT_INPUT_LIBRARY
+        return self._add_data(input_type, library, name)
+
+    def _add_data(self, input_type, data, name):
+        name_buffer = name.encode("utf8")
+        self._keep_alive += [data, name_buffer]
         try:
-            input_ptx = binding.CUjitInputType.CU_JIT_INPUT_PTX
             driver.cuLinkAddData(
-                self.handle, input_ptx, ptx, len(ptx), namebuf, 0, [], []
+                self.handle, input_type, data, len(data), name_buffer, 0, [], []
             )
         except CudaAPIError as e:
             raise LinkerError("%s\n%s" % (e, self.error_log))
 
+    def add_data(self, data, kind, name=None):
+        # We pass the name as **kwargs to ensure the default name for the input
+        # type is used if none is supplied
+        kws = {}
+        if name is not None:
+            kws["name"] = name
+
+        if kind == FILE_EXTENSION_MAP["cubin"]:
+            self.add_cubin(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["fatbin"]:
+            self.add_fatbin(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["a"]:
+            self.add_library(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["ptx"]:
+            self.add_ptx(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["o"]:
+            self.add_object(data, **kws)
+        elif kind == FILE_EXTENSION_MAP["ltoir"]:
+            raise LinkerError("CudaPythonLinker cannot link LTO-IR")
+        else:
+            raise LinkerError(f"Don't know how to link {kind}")
+
     def add_file(self, path, kind):
         pathbuf = path.encode("utf8")
         self._keep_alive.append(pathbuf)
@@ -3252,8 +3340,7 @@ class PyNvJitLinker(Linker):
 
     def add_file(self, path, kind):
         try:
-            with open(path, "rb") as f:
-                data = f.read()
+            data = cached_file_read(path, "rb")
         except FileNotFoundError:
             raise LinkerError(f"{path} not found")
 

{numba_cuda-0.10.1 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/cudadrv/linkable_code.py

@@ -16,16 +16,24 @@ class LinkableCode:
     :param teardown_callback: A function called just prior to the unloading of
                               a module that has this code object linked into
                               it.
+    :param nrt: If True, assume this object contains NRT function calls and
+                add NRT source code to the final link.
     """
 
     def __init__(
-        self, data, name=None, setup_callback=None, teardown_callback=None
+        self,
+        data,
+        name=None,
+        setup_callback=None,
+        teardown_callback=None,
+        nrt=False,
     ):
         if setup_callback and not callable(setup_callback):
             raise TypeError("setup_callback must be callable")
         if teardown_callback and not callable(teardown_callback):
             raise TypeError("teardown_callback must be callable")
 
+        self.nrt = nrt
         self._name = name
         self._data = data
         self.setup_callback = setup_callback
@@ -87,5 +95,5 @@ class Object(LinkableCode):
 class LTOIR(LinkableCode):
     """An LTOIR file in memory."""
 
-    kind = "ltoir"
+    kind = FILE_EXTENSION_MAP["ltoir"]
     default_name = "<unnamed-ltoir>"