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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: numba-cuda
-Version: 0.11.0
+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.11.0/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.11.0 → 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.11.0 → 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.11.0 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/codegen.py

@@ -70,6 +70,8 @@ class ExternalCodeLibrary(CodeLibrary):
         self._setup_functions = []
         self._teardown_functions = []
 
+        self.use_cooperative = False
+
     @property
     def modules(self):
         # There are no LLVM IR modules in an ExternalCodeLibrary
@@ -181,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:
@@ -352,6 +356,7 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         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):
@@ -442,6 +447,7 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
             nvvm_options=self._nvvm_options,
             needs_cudadevrt=self.needs_cudadevrt,
             nrt=nrt,
+            use_cooperative=self.use_cooperative,
         )
 
     @classmethod
@@ -458,6 +464,7 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
         nvvm_options,
         needs_cudadevrt,
         nrt,
+        use_cooperative,
     ):
         """
         Rebuild an instance.
@@ -472,6 +479,7 @@ 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:

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

@@ -797,7 +797,7 @@ def compile_ptx_for_current_device(
     )
 
 
-def declare_device_function(name, restype, argtypes, link):
+def declare_device_function(name, restype, argtypes, link, use_cooperative):
     from .descriptor import cuda_target
 
     typingctx = cuda_target.typing_context
@@ -816,6 +816,7 @@ def declare_device_function(name, restype, argtypes, link):
     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

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

@@ -423,7 +423,11 @@ _genfp16_binary_operator(operator.itruediv)
 def _resolve_wrapped_unary(fname):
     link = tuple()
     decl = declare_device_function(
-        f"__numba_wrapper_{fname}", types.float16, (types.float16,), link
+        f"__numba_wrapper_{fname}",
+        types.float16,
+        (types.float16,),
+        link,
+        use_cooperative=False,
     )
     return types.Function(decl)
 
@@ -438,6 +442,7 @@ def _resolve_wrapped_binary(fname):
             types.float16,
         ),
         link,
+        use_cooperative=False,
     )
     return types.Function(decl)
 

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

@@ -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:

{numba_cuda-0.11.0 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/cudadrv/nvrtc.py

@@ -6,13 +6,21 @@ from numba.cuda.cudadrv.error import (
     NvrtcCompilationError,
     NvrtcSupportError,
 )
+from numba import config
 from numba.cuda.cuda_paths import get_cuda_paths
+from numba.cuda.utils import _readenv
 
 import functools
 import os
 import threading
 import warnings
 
+NVRTC_EXTRA_SEARCH_PATHS = _readenv(
+    "NUMBA_CUDA_NVRTC_EXTRA_SEARCH_PATHS", str, ""
+) or getattr(config, "NUMBA_CUDA_NVRTC_EXTRA_SEARCH_PATHS", "")
+if not hasattr(config, "NUMBA_CUDA_NVRTC_EXTRA_SEARCH_PATHS"):
+    config.CUDA_NVRTC_EXTRA_SEARCH_PATHS = NVRTC_EXTRA_SEARCH_PATHS
+
 # Opaque handle for compilation unit
 nvrtc_program = c_void_p
 
@@ -383,10 +391,24 @@ def compile(src, name, cc, ltoir=False):
     else:
         numba_include = f"-I{os.path.join(numba_cuda_path, 'include', '12')}"
 
+    if config.CUDA_NVRTC_EXTRA_SEARCH_PATHS:
+        extra_search_paths = config.CUDA_NVRTC_EXTRA_SEARCH_PATHS.split(":")
+        extra_includes = [f"-I{p}" for p in extra_search_paths]
+    else:
+        extra_includes = []
+
     nrt_path = os.path.join(numba_cuda_path, "runtime")
     nrt_include = f"-I{nrt_path}"
 
-    options = [arch, numba_include, *cuda_include, nrt_include, "-rdc", "true"]
+    options = [
+        arch,
+        numba_include,
+        *cuda_include,
+        nrt_include,
+        *extra_includes,
+        "-rdc",
+        "true",
+    ]
 
     if ltoir:
         options.append("-dlto")

{numba_cuda-0.11.0 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/debuginfo.py

@@ -59,6 +59,33 @@ class CUDADIBuilder(DIBuilder):
         # For other cases, use upstream Numba implementation
         return super()._var_type(lltype, size, datamodel=datamodel)
 
+    def _di_subroutine_type(self, line, function, argmap):
+        # The function call conv needs encoding.
+        llfunc = function
+        md = []
+
+        # Create metadata type for return value
+        if len(llfunc.args) > 0:
+            lltype = llfunc.args[0].type
+            size = self.cgctx.get_abi_sizeof(lltype)
+            mdtype = self._var_type(lltype, size, datamodel=None)
+            md.append(mdtype)
+
+        # Create metadata type for arguments
+        for idx, (name, nbtype) in enumerate(argmap.items()):
+            datamodel = self.cgctx.data_model_manager[nbtype]
+            lltype = self.cgctx.get_value_type(nbtype)
+            size = self.cgctx.get_abi_sizeof(lltype)
+            mdtype = self._var_type(lltype, size, datamodel=datamodel)
+            md.append(mdtype)
+
+        return self.module.add_debug_info(
+            "DISubroutineType",
+            {
+                "types": self.module.add_metadata(md),
+            },
+        )
+
     def mark_variable(
         self,
         builder,

{numba_cuda-0.11.0 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/decorators.py

@@ -229,7 +229,7 @@ def jit(
                 return disp
 
 
-def declare_device(name, sig, link=None):
+def declare_device(name, sig, link=None, use_cooperative=False):
     """
     Declare the signature of a foreign function. Returns a descriptor that can
     be used to call the function from a Python kernel.
@@ -238,6 +238,7 @@ def declare_device(name, sig, link=None):
     :type name: str
     :param sig: The Numba signature of the function.
     :param link: External code to link when calling the function.
+    :param use_cooperative: External code requires cooperative launch.
     """
     if link is None:
         link = tuple()
@@ -250,6 +251,8 @@ def declare_device(name, sig, link=None):
         msg = "Return type must be provided for device declarations"
         raise TypeError(msg)
 
-    template = declare_device_function(name, restype, argtypes, link)
+    template = declare_device_function(
+        name, restype, argtypes, link, use_cooperative
+    )
 
     return template.key

{numba_cuda-0.11.0 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/dispatcher.py

@@ -151,8 +151,8 @@ class _Kernel(serialize.ReduceMixin):
 
         asm = lib.get_asm_str()
 
-        # A kernel needs cooperative launch if grid_sync is being used.
-        self.cooperative = "cudaCGGetIntrinsicHandle" in asm
+        # The code library contains functions that require cooperative launch.
+        self.cooperative = lib.use_cooperative
         # We need to link against cudadevrt if grid sync is being used.
         if self.cooperative:
             lib.needs_cudadevrt = True

{numba_cuda-0.11.0 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/target.py

@@ -290,7 +290,16 @@ class CUDATargetContext(BaseContext):
 
 
 class CUDACallConv(MinimalCallConv):
-    pass
+    def decorate_function(self, fn, args, fe_argtypes, noalias=False):
+        """
+        Set names and attributes of function arguments.
+        """
+        assert not noalias
+        arginfo = self._get_arg_packer(fe_argtypes)
+        # Do not prefix "arg." on argument name, so that nvvm compiler
+        # can track debug info of argument more accurately
+        arginfo.assign_names(self.get_arguments(fn), args)
+        fn.args[0].name = ".ret"
 
 
 class CUDACABICallConv(BaseCallConv):

{numba_cuda-0.11.0 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/tests/cudapy/cache_usecases.py

@@ -203,18 +203,6 @@ def simple_usecase_kernel(r, x):
 simple_usecase_caller = CUDAUseCase(simple_usecase_kernel)
 
 
-# Usecase with cooperative groups
-
-
-@cuda.jit(cache=True)
-def cg_usecase_kernel(r, x):
-    grid = cuda.cg.this_grid()
-    grid.sync()
-
-
-cg_usecase = CUDAUseCase(cg_usecase_kernel)
-
-
 class _TestModule(CUDATestCase):
     """
     Tests for functionality of this module's functions.

numba_cuda-0.12.1/numba_cuda/numba/cuda/tests/cudapy/cg_cache_usecases.py

@@ -0,0 +1,33 @@
+from numba import cuda
+from numba.cuda.testing import CUDATestCase
+import sys
+
+from numba.cuda.tests.cudapy.cache_usecases import CUDAUseCase
+
+
+# Usecase with cooperative groups
+
+
+@cuda.jit(cache=True)
+def cg_usecase_kernel(r, x):
+    grid = cuda.cg.this_grid()
+    grid.sync()
+
+
+cg_usecase = CUDAUseCase(cg_usecase_kernel)
+
+
+class _TestModule(CUDATestCase):
+    """
+    Tests for functionality of this module's functions.
+    Note this does not define any "test_*" method, instead check_module()
+    should be called by hand.
+    """
+
+    def check_module(self, mod):
+        mod.cg_usecase(0)
+
+
+def self_test():
+    mod = sys.modules[__name__]
+    _TestModule().check_module(mod)

numba_cuda-0.12.1/numba_cuda/numba/cuda/tests/cudapy/test_bfloat16.py

@@ -0,0 +1,55 @@
+from numba import cuda, float32
+from numba.cuda.bf16 import bfloat16
+from numba.cuda.testing import CUDATestCase
+
+import math
+
+
+class TestBfloat16HighLevelBindings(CUDATestCase):
+    def skip_unsupported(self):
+        if not cuda.is_bfloat16_supported():
+            self.skipTest(
+                "bfloat16 requires compute capability 8.0+ and CUDA version>= 12.0"
+            )
+
+    def test_use_type_in_kernel(self):
+        self.skip_unsupported()
+
+        @cuda.jit
+        def kernel():
+            bfloat16(3.14)
+
+        kernel[1, 1]()
+
+    def test_math_bindings(self):
+        self.skip_unsupported()
+        functions = [
+            math.trunc,
+            math.ceil,
+            math.floor,
+            math.sqrt,
+            math.log,
+            math.log10,
+            math.cos,
+            math.sin,
+            math.tanh,
+            math.exp,
+            math.exp2,
+        ]
+
+        for f in functions:
+            with self.subTest(func=f):
+
+                @cuda.jit
+                def kernel(arr):
+                    x = bfloat16(3.14)
+                    y = f(x)
+                    arr[0] = float32(y)
+
+                arr = cuda.device_array((1,), dtype="float32")
+                kernel[1, 1](arr)
+
+                if f in (math.exp, math.exp2):
+                    self.assertAlmostEqual(arr[0], f(3.14), delta=1e-1)
+                else:
+                    self.assertAlmostEqual(arr[0], f(3.14), delta=1e-2)

{numba_cuda-0.11.0 → numba_cuda-0.12.1}/numba_cuda/numba/cuda/tests/cudapy/test_bfloat16_bindings.py

@@ -5,7 +5,7 @@ import numpy as np
 from numba import int16, int32, int64, uint16, uint32, uint64, float32, float64
 from numba.types import float16
 
-from numba.cuda.cuda_bf16 import (
+from numba.cuda._internal.cuda_bf16 import (
     nv_bfloat16,
     htrunc,
     hceil,
@@ -22,21 +22,23 @@ from numba.cuda.cuda_bf16 import (
     hexp,
     hexp2,
     hexp10,
+    htanh,
+    htanh_approx,
 )
 
-from numba.cuda.cudadrv.runtime import get_version
-
-cuda_version = get_version()
-
 dtypes = [int16, int32, int64, uint16, uint32, uint64, float32]
 
 
-@unittest.skipIf(
-    (cuda.get_current_device().compute_capability < (8, 0)),
-    "bfloat16 requires compute capability 8.0+",
-)
 class Bfloat16Test(CUDATestCase):
+    def skip_unsupported(self):
+        if not cuda.is_bfloat16_supported():
+            self.skipTest(
+                "bfloat16 requires compute capability 8.0+ and CUDA version>= 12.0"
+            )
+
     def test_ctor(self):
+        self.skip_unsupported()
+
         @cuda.jit
         def simple_kernel():
             a = nv_bfloat16(float64(1.0))  # noqa: F841
@@ -47,18 +49,13 @@ class Bfloat16Test(CUDATestCase):
             f = nv_bfloat16(uint16(6))  # noqa: F841
             g = nv_bfloat16(uint32(7))  # noqa: F841
             h = nv_bfloat16(uint64(8))  # noqa: F841
+            i = nv_bfloat16(float16(9))  # noqa: F841
 
         simple_kernel[1, 1]()
 
-        if cuda_version >= (12, 0):
-
-            @cuda.jit
-            def simple_kernel_fp16():
-                i = nv_bfloat16(float16(9))  # noqa: F841
-
-            simple_kernel_fp16[1, 1]()
-
     def test_casts(self):
+        self.skip_unsupported()
+
         @cuda.jit
         def simple_kernel(b, c, d, e, f, g, h):
             a = nv_bfloat16(3.14)
@@ -90,6 +87,7 @@ class Bfloat16Test(CUDATestCase):
         assert h[0] == 3
 
     def test_ctor_cast_loop(self):
+        self.skip_unsupported()
         for dtype in dtypes:
             with self.subTest(dtype=dtype):
 
@@ -106,6 +104,8 @@ class Bfloat16Test(CUDATestCase):
                     assert a[0] == 3
 
     def test_arithmetic(self):
+        self.skip_unsupported()
+
         @cuda.jit
         def simple_kernel(arith, logic):
             # Binary Arithmetic Operators
@@ -175,6 +175,8 @@ class Bfloat16Test(CUDATestCase):
         )
 
     def test_math_func(self):
+        self.skip_unsupported()
+
         @cuda.jit
         def simple_kernel(a):
             x = nv_bfloat16(3.14)
@@ -191,16 +193,18 @@ class Bfloat16Test(CUDATestCase):
             a[9] = float32(hlog10(x))
             a[10] = float32(hcos(x))
             a[11] = float32(hsin(x))
-            a[12] = float32(hexp(x))
-            a[13] = float32(hexp2(x))
-            a[14] = float32(hexp10(x))
+            a[12] = float32(htanh(x))
+            a[13] = float32(htanh_approx(x))
+            a[14] = float32(hexp(x))
+            a[15] = float32(hexp2(x))
+            a[16] = float32(hexp10(x))
 
-        a = np.zeros(15, dtype=np.float32)
+        a = np.zeros(17, dtype=np.float32)
         simple_kernel[1, 1](a)
 
         x = 3.14
         np.testing.assert_allclose(
-            a[:12],
+            a[:14],
             [
                 np.trunc(x),
                 np.ceil(x),
@@ -214,15 +218,19 @@ class Bfloat16Test(CUDATestCase):
                 np.log10(x),
                 np.cos(x),
                 np.sin(x),
+                np.tanh(x),
+                np.tanh(x),
             ],
             atol=1e-2,
         )
 
         np.testing.assert_allclose(
-            a[12:], [np.exp(x), np.exp2(x), np.power(10, x)], atol=1e2
+            a[14:], [np.exp(x), np.exp2(x), np.power(10, x)], atol=1e2
         )
 
     def test_check_bfloat16_type(self):
+        self.skip_unsupported()
+
         @cuda.jit
         def kernel(arr):
             x = nv_bfloat16(3.14)
@@ -237,6 +245,8 @@ class Bfloat16Test(CUDATestCase):
         np.testing.assert_allclose(arr, [3.14], atol=1e-2)
 
     def test_use_within_device_func(self):
+        self.skip_unsupported()
+
         @cuda.jit(device=True)
         def add_bf16(a, b):
             return a + b
@@ -252,6 +262,22 @@ class Bfloat16Test(CUDATestCase):
 
         np.testing.assert_allclose(arr, [8], atol=1e-2)
 
+    def test_use_binding_inside_dfunc(self):
+        @cuda.jit(device=True)
+        def f(arr):
+            pi = nv_bfloat16(3.14)
+            three = htrunc(pi)
+            arr[0] = float32(three)
+
+        @cuda.jit
+        def kernel(arr):
+            f(arr)
+
+        arr = np.zeros(1, np.float32)
+        kernel[1, 1](arr)
+
+        np.testing.assert_allclose(arr, [3], atol=1e-2)
+
 
 if __name__ == "__main__":
     unittest.main()