numba-cuda 0.19.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

numba_cuda/numba/cuda/initialize.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+
 def initialize_all():
     # Import models to register them with the data model manager
     import numba.cuda.models  # noqa: F401

numba_cuda/numba/cuda/intrinsic_wrapper.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from .decorators import jit
 import numba
 

numba_cuda/numba/cuda/intrinsics.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from llvmlite import ir
 
 from numba import cuda, types

numba_cuda/numba/cuda/itanium_mangler.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 """
 Itanium CXX ABI Mangler
 

numba_cuda/numba/cuda/kernels/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause

numba_cuda/numba/cuda/kernels/reduction.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 """
 A library written in CUDA Python for generating reduction kernels
 """

numba_cuda/numba/cuda/kernels/transpose.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from numba import cuda
 from numba.cuda.cudadrv.driver import driver
 import math

numba_cuda/numba/cuda/libdevice.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+
 def abs(x):
     """
     See https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_abs.html

numba_cuda/numba/cuda/libdevicedecl.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from numba.cuda import libdevice, libdevicefuncs
 from numba.cuda.typing.templates import ConcreteTemplate, Registry
 

numba_cuda/numba/cuda/libdevicefuncs.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from collections import namedtuple
 from textwrap import indent
 

numba_cuda/numba/cuda/libdeviceimpl.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from llvmlite import ir
 from numba.core import types
 from numba.cuda import cgutils

numba_cuda/numba/cuda/locks.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from threading import Lock
 from functools import wraps
 

numba_cuda/numba/cuda/lowering.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from collections import namedtuple, defaultdict
 import operator
 import warnings
@@ -7,18 +10,14 @@ from llvmlite import ir as llvm_ir
 
 from numba.core import (
     typing,
-    utils,
     types,
     ir,
-    debuginfo,
-    funcdesc,
     generators,
-    config,
-    cgutils,
     removerefctpass,
-    targetconfig,
 )
-from numba.cuda.core import ir_utils
+from numba.cuda import debuginfo, cgutils, utils
+from numba.cuda.core import ir_utils, targetconfig, funcdesc, config
+
 from numba.core.errors import (
     LoweringError,
     new_error_context,
@@ -27,8 +26,8 @@ from numba.core.errors import (
     UnsupportedError,
     NumbaDebugInfoWarning,
 )
-from numba.core.funcdesc import default_mangler
-from numba.core.environment import Environment
+from numba.cuda.core.funcdesc import default_mangler
+from numba.cuda.core.environment import Environment
 from numba.core.analysis import compute_use_defs, must_use_alloca
 from numba.misc.firstlinefinder import get_func_body_first_lineno
 from numba import version_info
@@ -463,7 +462,7 @@ class Lower(BaseLower):
         self._blk_local_varmap = {}
 
     def pre_block(self, block):
-        from numba.core.unsafe import eh
+        from numba.cuda.core.unsafe import eh
 
         super(Lower, self).pre_block(block)
         self._cur_ir_block = block
@@ -1026,9 +1025,6 @@ class Lower(BaseLower):
         elif isinstance(fnty, types.RecursiveCall):
             res = self._lower_call_RecursiveCall(fnty, expr, signature)
 
-        elif isinstance(fnty, types.FunctionType):
-            res = self._lower_call_FunctionType(fnty, expr, signature)
-
         else:
             res = self._lower_call_normal(fnty, expr, signature)
 
@@ -1049,7 +1045,7 @@ class Lower(BaseLower):
         )
 
     def _lower_call_ObjModeDispatcher(self, fnty, expr, signature):
-        from numba.core.pythonapi import ObjModeUtils
+        from numba.cuda.core.pythonapi import ObjModeUtils
 
         self.init_pyapi()
         # Acquire the GIL
@@ -1226,136 +1222,6 @@ class Lower(BaseLower):
             )
         return res
 
-    def _lower_call_FunctionType(self, fnty, expr, signature):
-        self.debug_print("# calling first-class function type")
-        sig = types.unliteral(signature)
-        if not fnty.check_signature(signature):
-            # value dependent polymorphism?
-            raise UnsupportedError(
-                f"mismatch of function types:"
-                f" expected {fnty} but got {types.FunctionType(sig)}"
-            )
-        argvals = self.fold_call_args(
-            fnty,
-            sig,
-            expr.args,
-            expr.vararg,
-            expr.kws,
-        )
-        return self.__call_first_class_function_pointer(
-            fnty.ftype,
-            expr.func.name,
-            sig,
-            argvals,
-        )
-
-    def __call_first_class_function_pointer(self, ftype, fname, sig, argvals):
-        """
-        Calls a first-class function pointer.
-
-        This function is responsible for calling a first-class function pointer,
-        which can either be a JIT-compiled function or a Python function. It
-        determines if a JIT address is available, and if so, calls the function
-        using the JIT address. Otherwise, it calls the function using a function
-        pointer obtained from the `__get_first_class_function_pointer` method.
-
-        Args:
-            ftype: The type of the function.
-            fname: The name of the function.
-            sig: The signature of the function.
-            argvals: The argument values to pass to the function.
-
-        Returns:
-            The result of calling the function.
-        """
-        context = self.context
-        builder = self.builder
-        # Determine if jit address is available
-        fstruct = self.loadvar(fname)
-        struct = cgutils.create_struct_proxy(self.typeof(fname))(
-            context, builder, value=fstruct
-        )
-        jit_addr = struct.jit_addr
-        jit_addr.name = f"jit_addr_of_{fname}"
-
-        ctx = context
-        res_slot = cgutils.alloca_once(
-            builder, ctx.get_value_type(sig.return_type)
-        )
-
-        if_jit_addr_is_null = builder.if_else(
-            cgutils.is_null(builder, jit_addr), likely=False
-        )
-        with if_jit_addr_is_null as (then, orelse):
-            with then:
-                func_ptr = self.__get_first_class_function_pointer(
-                    ftype, fname, sig
-                )
-                res = builder.call(func_ptr, argvals)
-                builder.store(res, res_slot)
-
-            with orelse:
-                llty = ctx.call_conv.get_function_type(
-                    sig.return_type, sig.args
-                ).as_pointer()
-                func_ptr = builder.bitcast(jit_addr, llty)
-                # call
-                status, res = ctx.call_conv.call_function(
-                    builder, func_ptr, sig.return_type, sig.args, argvals
-                )
-                with cgutils.if_unlikely(builder, status.is_error):
-                    context.call_conv.return_status_propagate(builder, status)
-                builder.store(res, res_slot)
-        return builder.load(res_slot)
-
-    def __get_first_class_function_pointer(self, ftype, fname, sig):
-        from numba.experimental.function_type import lower_get_wrapper_address
-
-        llty = self.context.get_value_type(ftype)
-        fstruct = self.loadvar(fname)
-        addr = self.builder.extract_value(
-            fstruct, 0, name="addr_of_%s" % (fname)
-        )
-
-        fptr = cgutils.alloca_once(
-            self.builder, llty, name="fptr_of_%s" % (fname)
-        )
-        with self.builder.if_else(
-            cgutils.is_null(self.builder, addr), likely=False
-        ) as (then, orelse):
-            with then:
-                self.init_pyapi()
-                # Acquire the GIL
-                gil_state = self.pyapi.gil_ensure()
-                pyaddr = self.builder.extract_value(
-                    fstruct, 1, name="pyaddr_of_%s" % (fname)
-                )
-                # try to recover the function address, see
-                # test_zero_address BadToGood example in
-                # test_function_type.py
-                addr1 = lower_get_wrapper_address(
-                    self.context,
-                    self.builder,
-                    pyaddr,
-                    sig,
-                    failure_mode="ignore",
-                )
-                with self.builder.if_then(
-                    cgutils.is_null(self.builder, addr1), likely=False
-                ):
-                    self.return_exception(
-                        RuntimeError,
-                        exc_args=(f"{ftype} function address is null",),
-                        loc=self.loc,
-                    )
-                addr2 = self.pyapi.long_as_voidptr(addr1)
-                self.builder.store(self.builder.bitcast(addr2, llty), fptr)
-                self.pyapi.decref(addr1)
-                self.pyapi.gil_release(gil_state)
-            with orelse:
-                self.builder.store(self.builder.bitcast(addr, llty), fptr)
-        return self.builder.load(fptr)
-
     def _lower_call_normal(self, fnty, expr, signature):
         # Normal function resolution
         self.debug_print("# calling normal function: {0}".format(fnty))
@@ -1841,8 +1707,7 @@ class CUDALower(Lower):
             int_type = (llvm_ir.IntType,)
             real_type = llvm_ir.FloatType, llvm_ir.DoubleType
             if isinstance(lltype, int_type + real_type):
-                index = name.find(".")
-                src_name = name[:index] if index > 0 else name
+                src_name = name.split(".")[0]
                 if src_name in self.poly_var_typ_map:
                     # Do not emit debug value on polymorphic type var
                     return
@@ -1869,6 +1734,9 @@ class CUDALower(Lower):
 
         self.poly_var_typ_map = {}
         self.poly_var_loc_map = {}
+        self.poly_var_set = set()
+        self.poly_cleaned = False
+        self.lastblk = max(self.blocks.keys())
 
         # When debug info is enabled, walk through function body and mark
         # variables with polymorphic types.
@@ -1880,8 +1748,7 @@ class CUDALower(Lower):
                     if x.target.name.startswith("$"):
                         continue
                     ssa_name = x.target.name
-                    index = ssa_name.find(".")
-                    src_name = ssa_name[:index] if index > 0 else ssa_name
+                    src_name = ssa_name.split(".")[0]
                     # Check all the multi-versioned targets
                     if len(x.target.versioned_names) > 0:
                         fetype = self.typeof(ssa_name)
@@ -1902,12 +1769,12 @@ class CUDALower(Lower):
         """
         # If the name is not handled yet and a store is needed
         if name not in self.varmap and self.store_var_needed(name):
-            index = name.find(".")
-            src_name = name[:index] if index > 0 else name
+            src_name = name.split(".")[0]
             if src_name in self.poly_var_typ_map:
-                dtype = types.UnionType(self.poly_var_typ_map[src_name])
-                datamodel = self.context.data_model_manager[dtype]
+                self.poly_var_set.add(name)
                 if src_name not in self.poly_var_loc_map:
+                    dtype = types.UnionType(self.poly_var_typ_map[src_name])
+                    datamodel = self.context.data_model_manager[dtype]
                     # UnionType has sorted set of types, max at last index
                     maxsizetype = dtype.types[-1]
                     # Create a single element aggregate type
@@ -1916,13 +1783,7 @@ class CUDALower(Lower):
                     ptr = self.alloca_lltype(src_name, lltype, datamodel)
                     # save the location of the union type for polymorphic var
                     self.poly_var_loc_map[src_name] = ptr
-                # Any member of this union type shoud type cast ptr to fetype
-                lltype = self.context.get_value_type(fetype)
-                castptr = self.builder.bitcast(
-                    self.poly_var_loc_map[src_name], llvm_ir.PointerType(lltype)
-                )
-                # Remember the pointer
-                self.varmap[name] = castptr
+                return
 
         super()._alloca_var(name, fetype)
 
@@ -1936,6 +1797,45 @@ class CUDALower(Lower):
             or self._disable_sroa_like_opt
         )
 
+    def delvar(self, name):
+        """
+        Delete the given variable.
+        """
+        if name in self.poly_var_set:
+            fetype = self.typeof(name)
+            src_name = name.split(".")[0]
+            ptr = self.poly_var_loc_map[src_name]
+            self.decref(fetype, self.builder.load(ptr))
+            if (
+                self._cur_ir_block == self.blocks[self.lastblk]
+                and not self.poly_cleaned
+            ):
+                # Zero-fill the debug union for polymorphic only
+                # at the last block
+                for v in self.poly_var_loc_map.values():
+                    self.builder.store(
+                        llvm_ir.Constant(v.type.pointee, None), v
+                    )
+                    self.poly_cleaned = True
+            return
+
+        super().delvar(name)
+
+    def getvar(self, name):
+        """
+        Get a pointer to the given variable's slot.
+        """
+        if name in self.poly_var_set:
+            src_name = name.split(".")[0]
+            fetype = self.typeof(name)
+            lltype = self.context.get_value_type(fetype)
+            castptr = self.builder.bitcast(
+                self.poly_var_loc_map[src_name], llvm_ir.PointerType(lltype)
+            )
+            return castptr
+        else:
+            return super().getvar(name)
+
 
 def _lit_or_omitted(value):
     """Returns a Literal instance if the type of value is supported;

numba_cuda/numba/cuda/mathimpl.py

@@ -1,11 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 import math
 import operator
 from llvmlite import ir
-from numba.core import types, typing, targetconfig
+from numba.core import types, typing
 from numba.cuda import cgutils
 from numba.core.imputils import Registry
 from numba.types import float32, float64, int64, uint64
 from numba.cuda import libdevice
+from numba.cuda.core import targetconfig
 
 registry = Registry()
 lower = registry.lower

numba_cuda/numba/cuda/memory_management/__init__.py

@@ -1 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 from numba.cuda.memory_management.nrt import rtsys  # noqa: F401

numba_cuda/numba/cuda/memory_management/memsys.cu

@@ -1,3 +1,8 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
 #include "memsys.cuh"
 
 __device__ size_t memsys_size = sizeof(NRT_MemSys);

numba_cuda/numba/cuda/memory_management/memsys.cuh

@@ -1,3 +1,8 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
 #include <cuda/atomic>
 
 // Globally needed variables

numba_cuda/numba/cuda/memory_management/nrt.cu

@@ -1,3 +1,8 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
 #ifndef _NRT_H
 #define _NRT_H
 

numba_cuda/numba/cuda/memory_management/nrt.cuh

@@ -1,3 +1,8 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
 #include <cuda/atomic>
 
 typedef void (*NRT_dtor_function)(void* ptr, size_t size, void* info);

numba_cuda/numba/cuda/memory_management/nrt.py

@@ -1,36 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 import ctypes
 import os
 from functools import wraps
 import numpy as np
 
-from numba import cuda, config
+
+from numba import cuda, types
+from numba.cuda import config
+
 from numba.core.runtime.nrt import _nrt_mstats
 from numba.cuda.cudadrv.driver import (
     _Linker,
     driver,
     launch_kernel,
     USE_NV_BINDING,
+    _have_nvjitlink,
 )
 from numba.cuda.cudadrv import devices
 from numba.cuda.api import get_current_device
-from numba.cuda.utils import _readenv, cached_file_read
+from numba.cuda.utils import cached_file_read
 from numba.cuda.cudadrv.linkable_code import CUSource
+from numba.cuda.typing.templates import signature
 
-
-# Check environment variable or config for NRT statistics enablement
-NRT_STATS = _readenv("NUMBA_CUDA_NRT_STATS", bool, False) or getattr(
-    config, "NUMBA_CUDA_NRT_STATS", False
-)
-if not hasattr(config, "NUMBA_CUDA_NRT_STATS"):
-    config.CUDA_NRT_STATS = NRT_STATS
-
-
-# Check environment variable or config for NRT enablement
-ENABLE_NRT = _readenv("NUMBA_CUDA_ENABLE_NRT", bool, False) or getattr(
-    config, "NUMBA_CUDA_ENABLE_NRT", False
-)
-if not hasattr(config, "NUMBA_CUDA_ENABLE_NRT"):
-    config.CUDA_ENABLE_NRT = ENABLE_NRT
+from numba.core.extending import intrinsic, overload_classmethod
 
 
 def get_include():
@@ -38,6 +32,34 @@ def get_include():
     return os.path.dirname(os.path.abspath(__file__))
 
 
+# Provide an implementation of Array._allocate() for the CUDA target (used
+# internally by Numba when generating the allocation of an array)
+
+
+@intrinsic
+def intrin_alloc(typingctx, allocsize, align):
+    """Intrinsic to call into the allocator for Array"""
+
+    def codegen(context, builder, signature, args):
+        allocsize, align = args
+        meminfo = context.nrt.meminfo_alloc_aligned(builder, allocsize, align)
+        return meminfo
+
+    mip = types.MemInfoPointer(types.voidptr)  # return untyped pointer
+    sig = signature(mip, allocsize, align)
+    return sig, codegen
+
+
+@overload_classmethod(types.Array, "_allocate", target="CUDA")
+def _ol_array_allocate(cls, allocsize, align):
+    """Implements a Numba-only CUDA-target classmethod on the array type."""
+
+    def impl(cls, allocsize, align):
+        return intrin_alloc(allocsize, align)
+
+    return impl
+
+
 # Protect method to ensure NRT memory allocation and initialization
 def _alloc_init_guard(method):
     """
@@ -65,10 +87,18 @@ class _Runtime:
 
     def __init__(self):
         """Initialize memsys module and variable"""
+        self._reset()
+
+    def _reset(self):
+        """Reset to the uninitialized state"""
         self._memsys_module = None
         self._memsys = None
         self._initialized = False
 
+    def close(self):
+        """Close and reset"""
+        self._reset()
+
     def _compile_memsys_module(self):
         """
         Compile memsys.cu and create a module from it in the current context
@@ -80,7 +110,7 @@ class _Runtime:
         cc = get_current_device().compute_capability
 
         # Create a new linker instance and add the cu file
-        linker = _Linker.new(cc=cc)
+        linker = _Linker.new(cc=cc, lto=_have_nvjitlink())
         linker.add_cu_file(memsys_mod)
 
         # Complete the linker and create a module from it

numba_cuda/numba/cuda/misc/findlib.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+import os
+import re
+
+
+def get_lib_dirs():
+    """
+    Anaconda specific
+    """
+    if sys.platform == "win32":
+        # CUDA 12 puts in "bin" directory, whereas CUDA 13 puts in "bin\x64" directory
+        dirnames = [
+            os.path.join("Library", "bin"),
+            os.path.join("Library", "bin", "x64"),
+            os.path.join("Library", "nvvm", "bin"),
+            os.path.join("Library", "nvvm", "bin", "x64"),
+        ]
+    else:
+        dirnames = [
+            "lib",
+        ]
+    libdirs = [os.path.join(sys.prefix, x) for x in dirnames]
+    return libdirs
+
+
+DLLNAMEMAP = {
+    "linux": r"lib%(name)s\.so\.%(ver)s$",
+    "linux2": r"lib%(name)s\.so\.%(ver)s$",
+    "linux-static": r"lib%(name)s\.a$",
+    "darwin": r"lib%(name)s\.%(ver)s\.dylib$",
+    "win32": r"%(name)s%(ver)s\.dll$",
+    "win32-static": r"%(name)s\.lib$",
+    "bsd": r"lib%(name)s\.so\.%(ver)s$",
+}
+
+RE_VER = r"[0-9]*([_\.][0-9]+)*"
+
+
+def find_lib(libname, libdir=None, platform=None, static=False):
+    platform = platform or sys.platform
+    platform = "bsd" if "bsd" in platform else platform
+    if static:
+        platform = f"{platform}-static"
+    if platform not in DLLNAMEMAP:
+        # Return empty list if platform name is undefined.
+        # Not all platforms define their static library paths.
+        return []
+    pat = DLLNAMEMAP[platform] % {"name": libname, "ver": RE_VER}
+    regex = re.compile(pat)
+    return find_file(regex, libdir)
+
+
+def find_file(pat, libdir=None):
+    if libdir is None:
+        libdirs = get_lib_dirs()
+    elif isinstance(libdir, str):
+        libdirs = [
+            libdir,
+        ]
+    else:
+        libdirs = list(libdir)
+    files = []
+    for ldir in libdirs:
+        try:
+            entries = os.listdir(ldir)
+        except FileNotFoundError:
+            continue
+        candidates = [
+            os.path.join(ldir, ent) for ent in entries if pat.match(ent)
+        ]
+        files.extend([c for c in candidates if os.path.isfile(c)])
+    return files

numba_cuda/numba/cuda/models.py

@@ -1,11 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
 import functools
 
 from llvmlite import ir
 
 from numba.core.datamodel.registry import DataModelManager, register
+from numba.core.datamodel import PrimitiveModel
 from numba.core.extending import models
 from numba.core import types
-from numba.cuda.types import Dim3, GridGroup, CUDADispatcher
+from numba.cuda.types import Dim3, GridGroup, CUDADispatcher, Bfloat16
 
 
 cuda_data_manager = DataModelManager()
@@ -42,3 +46,10 @@ class FloatModel(models.PrimitiveModel):
 
 
 register_model(CUDADispatcher)(models.OpaqueModel)
+
+
+@register_model(Bfloat16)
+class _model___nv_bfloat16(PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        be_type = ir.IntType(16)
+        super(_model___nv_bfloat16, self).__init__(dmm, fe_type, be_type)