numba-cuda 0.22.0__cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/args.py

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Hints to wrap Kernel arguments to indicate how to manage host-device
+memory transfers before & after the kernel call.
+"""
+
+from numba.cuda.typing.typeof import typeof, Purpose
+
+
+class ArgHint:
+    def __init__(self, value):
+        self.value = value
+
+    def to_device(self, retr, stream=0):
+        """
+        :param stream: a stream to use when copying data
+        :param retr:
+            a list of clean-up work to do after the kernel's been run.
+            Append 0-arg lambdas to it!
+        :return: a value (usually an `DeviceNDArray`) to be passed to
+            the kernel
+        """
+
+    @property
+    def _numba_type_(self):
+        return typeof(self.value, Purpose.argument)
+
+
+class In(ArgHint):
+    def to_device(self, retr, stream=0):
+        from .cudadrv.devicearray import auto_device
+
+        devary, _ = auto_device(self.value, stream=stream)
+        # A dummy writeback functor to keep devary alive until the kernel
+        # is called.
+        retr.append(lambda: devary)
+        return devary
+
+
+class Out(ArgHint):
+    def to_device(self, retr, stream=0):
+        from .cudadrv.devicearray import auto_device
+
+        devary, conv = auto_device(self.value, copy=False, stream=stream)
+        if conv:
+            retr.append(lambda: devary.copy_to_host(self.value, stream=stream))
+        return devary
+
+
+class InOut(ArgHint):
+    def to_device(self, retr, stream=0):
+        from .cudadrv.devicearray import auto_device
+
+        devary, conv = auto_device(self.value, stream=stream)
+        if conv:
+            retr.append(lambda: devary.copy_to_host(self.value, stream=stream))
+        return devary
+
+
+def wrap_arg(value, default=InOut):
+    return value if isinstance(value, ArgHint) else default(value)
+
+
+__all__ = [
+    "In",
+    "Out",
+    "InOut",
+    "ArgHint",
+    "wrap_arg",
+]

numba_cuda/numba/cuda/bf16.py

@@ -0,0 +1,397 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+import sys
+
+from numba.cuda._internal.cuda_bf16 import (
+    typing_registry,
+    target_registry,
+    nv_bfloat16 as bfloat16,
+    # Arithmetic intrinsics
+    __habs as habs,
+    __hadd as hadd,
+    __hsub as hsub,
+    __hmul as hmul,
+    __hadd_rn as hadd_rn,
+    __hsub_rn as hsub_rn,
+    __hmul_rn as hmul_rn,
+    __hdiv as hdiv,
+    __hadd_sat as hadd_sat,
+    __hsub_sat as hsub_sat,
+    __hmul_sat as hmul_sat,
+    __hfma as hfma,
+    __hfma_sat as hfma_sat,
+    __hneg as hneg,
+    __hfma_relu as hfma_relu,
+    # Comparison intrinsics
+    __heq as heq,
+    __hne as hne,
+    __hge as hge,
+    __hgt as hgt,
+    __hle as hle,
+    __hlt as hlt,
+    __hmax as hmax,
+    __hmin as hmin,
+    __hmax_nan as hmax_nan,
+    __hmin_nan as hmin_nan,
+    __hisinf as hisinf,
+    __hisnan as hisnan,
+    # Unordered comparison intrinsics
+    __hequ as hequ,
+    __hneu as hneu,
+    __hgeu as hgeu,
+    __hgtu as hgtu,
+    __hleu as hleu,
+    __hltu as hltu,
+    # Precision conversion and data movement
+    # - floating-point family
+    __bfloat162float as bfloat162float,
+    __float2bfloat16 as float2bfloat16,
+    __double2bfloat16 as double2bfloat16,
+    __float2bfloat16_rn as float2bfloat16_rn,
+    __float2bfloat16_rz as float2bfloat16_rz,
+    __float2bfloat16_rd as float2bfloat16_rd,
+    __float2bfloat16_ru as float2bfloat16_ru,
+    # - char family
+    __bfloat162char_rz as bfloat162char_rz,
+    __bfloat162uchar_rz as bfloat162uchar_rz,
+    # - int family (signed 32-bit)
+    __int2bfloat16_rn as int2bfloat16_rn,
+    __int2bfloat16_rz as int2bfloat16_rz,
+    __int2bfloat16_rd as int2bfloat16_rd,
+    __int2bfloat16_ru as int2bfloat16_ru,
+    __bfloat162int_rn as bfloat162int_rn,
+    __bfloat162int_rz as bfloat162int_rz,
+    __bfloat162int_rd as bfloat162int_rd,
+    __bfloat162int_ru as bfloat162int_ru,
+    # - short family (signed 16-bit)
+    __short2bfloat16_rn as short2bfloat16_rn,
+    __short2bfloat16_rz as short2bfloat16_rz,
+    __short2bfloat16_rd as short2bfloat16_rd,
+    __short2bfloat16_ru as short2bfloat16_ru,
+    __bfloat162short_rn as bfloat162short_rn,
+    __bfloat162short_rz as bfloat162short_rz,
+    __bfloat162short_rd as bfloat162short_rd,
+    __bfloat162short_ru as bfloat162short_ru,
+    # - ushort family (unsigned 16-bit)
+    __ushort2bfloat16_rn as ushort2bfloat16_rn,
+    __ushort2bfloat16_rz as ushort2bfloat16_rz,
+    __ushort2bfloat16_rd as ushort2bfloat16_rd,
+    __ushort2bfloat16_ru as ushort2bfloat16_ru,
+    __bfloat162ushort_rn as bfloat162ushort_rn,
+    __bfloat162ushort_rz as bfloat162ushort_rz,
+    __bfloat162ushort_rd as bfloat162ushort_rd,
+    __bfloat162ushort_ru as bfloat162ushort_ru,
+    # - uint family (unsigned 32-bit)
+    __uint2bfloat16_rn as uint2bfloat16_rn,
+    __uint2bfloat16_rz as uint2bfloat16_rz,
+    __uint2bfloat16_rd as uint2bfloat16_rd,
+    __uint2bfloat16_ru as uint2bfloat16_ru,
+    __bfloat162uint_rn as bfloat162uint_rn,
+    __bfloat162uint_rz as bfloat162uint_rz,
+    __bfloat162uint_rd as bfloat162uint_rd,
+    __bfloat162uint_ru as bfloat162uint_ru,
+    # - ll family (signed 64-bit)
+    __ll2bfloat16_rn as ll2bfloat16_rn,
+    __ll2bfloat16_rz as ll2bfloat16_rz,
+    __ll2bfloat16_rd as ll2bfloat16_rd,
+    __ll2bfloat16_ru as ll2bfloat16_ru,
+    __bfloat162ll_rn as bfloat162ll_rn,
+    __bfloat162ll_rz as bfloat162ll_rz,
+    __bfloat162ll_rd as bfloat162ll_rd,
+    __bfloat162ll_ru as bfloat162ll_ru,
+    # - ull family (unsigned 64-bit)
+    __ull2bfloat16_rn as ull2bfloat16_rn,
+    __ull2bfloat16_rz as ull2bfloat16_rz,
+    __ull2bfloat16_rd as ull2bfloat16_rd,
+    __ull2bfloat16_ru as ull2bfloat16_ru,
+    __bfloat162ull_rn as bfloat162ull_rn,
+    __bfloat162ull_rz as bfloat162ull_rz,
+    __bfloat162ull_rd as bfloat162ull_rd,
+    __bfloat162ull_ru as bfloat162ull_ru,
+    # - bit reinterpret casts
+    __bfloat16_as_short as bfloat16_as_short,
+    __bfloat16_as_ushort as bfloat16_as_ushort,
+    __short_as_bfloat16 as short_as_bfloat16,
+    __ushort_as_bfloat16 as ushort_as_bfloat16,
+    htrunc,
+    hceil,
+    hfloor,
+    hrint,
+    hsqrt,
+    hrsqrt,
+    hrcp,
+    hlog,
+    hlog2,
+    hlog10,
+    hcos,
+    hsin,
+    hexp,
+    hexp2,
+    hexp10,
+    htanh,
+    htanh_approx,
+)
+from numba.cuda.extending import overload
+
+import math
+
+
+def _make_unary(a, func):
+    if a == 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)
+
+
+if sys.version_info >= (3, 11):
+
+    @overload(math.exp2, target="cuda")
+    def exp2_ol(a):
+        return _make_unary(a, hexp2)
+
+
+## Public aliases using Numba/Numpy-style type names
+# Floating-point
+float32_to_bfloat16 = float2bfloat16
+float64_to_bfloat16 = double2bfloat16
+bfloat16_to_float32 = bfloat162float
+float32_to_bfloat16_rn = float2bfloat16_rn
+float32_to_bfloat16_rz = float2bfloat16_rz
+float32_to_bfloat16_rd = float2bfloat16_rd
+float32_to_bfloat16_ru = float2bfloat16_ru
+
+# Char (8-bit)
+bfloat16_to_int8_rz = bfloat162char_rz
+bfloat16_to_uint8_rz = bfloat162uchar_rz
+
+# Int16 / UInt16
+int16_to_bfloat16_rn = short2bfloat16_rn
+int16_to_bfloat16_rz = short2bfloat16_rz
+int16_to_bfloat16_rd = short2bfloat16_rd
+int16_to_bfloat16_ru = short2bfloat16_ru
+bfloat16_to_int16_rn = bfloat162short_rn
+bfloat16_to_int16_rz = bfloat162short_rz
+bfloat16_to_int16_rd = bfloat162short_rd
+bfloat16_to_int16_ru = bfloat162short_ru
+
+uint16_to_bfloat16_rn = ushort2bfloat16_rn
+uint16_to_bfloat16_rz = ushort2bfloat16_rz
+uint16_to_bfloat16_rd = ushort2bfloat16_rd
+uint16_to_bfloat16_ru = ushort2bfloat16_ru
+bfloat16_to_uint16_rn = bfloat162ushort_rn
+bfloat16_to_uint16_rz = bfloat162ushort_rz
+bfloat16_to_uint16_rd = bfloat162ushort_rd
+bfloat16_to_uint16_ru = bfloat162ushort_ru
+
+# Int32 / UInt32
+int32_to_bfloat16_rn = int2bfloat16_rn
+int32_to_bfloat16_rz = int2bfloat16_rz
+int32_to_bfloat16_rd = int2bfloat16_rd
+int32_to_bfloat16_ru = int2bfloat16_ru
+bfloat16_to_int32_rn = bfloat162int_rn
+bfloat16_to_int32_rz = bfloat162int_rz
+bfloat16_to_int32_rd = bfloat162int_rd
+bfloat16_to_int32_ru = bfloat162int_ru
+
+uint32_to_bfloat16_rn = uint2bfloat16_rn
+uint32_to_bfloat16_rz = uint2bfloat16_rz
+uint32_to_bfloat16_rd = uint2bfloat16_rd
+uint32_to_bfloat16_ru = uint2bfloat16_ru
+bfloat16_to_uint32_rn = bfloat162uint_rn
+bfloat16_to_uint32_rz = bfloat162uint_rz
+bfloat16_to_uint32_rd = bfloat162uint_rd
+bfloat16_to_uint32_ru = bfloat162uint_ru
+
+# Int64 / UInt64
+int64_to_bfloat16_rn = ll2bfloat16_rn
+int64_to_bfloat16_rz = ll2bfloat16_rz
+int64_to_bfloat16_rd = ll2bfloat16_rd
+int64_to_bfloat16_ru = ll2bfloat16_ru
+bfloat16_to_int64_rn = bfloat162ll_rn
+bfloat16_to_int64_rz = bfloat162ll_rz
+bfloat16_to_int64_rd = bfloat162ll_rd
+bfloat16_to_int64_ru = bfloat162ll_ru
+
+uint64_to_bfloat16_rn = ull2bfloat16_rn
+uint64_to_bfloat16_rz = ull2bfloat16_rz
+uint64_to_bfloat16_rd = ull2bfloat16_rd
+uint64_to_bfloat16_ru = ull2bfloat16_ru
+bfloat16_to_uint64_rn = bfloat162ull_rn
+bfloat16_to_uint64_rz = bfloat162ull_rz
+bfloat16_to_uint64_rd = bfloat162ull_rd
+bfloat16_to_uint64_ru = bfloat162ull_ru
+
+# Bit reinterpret casts
+bfloat16_as_int16 = bfloat16_as_short
+bfloat16_as_uint16 = bfloat16_as_ushort
+int16_as_bfloat16 = short_as_bfloat16
+uint16_as_bfloat16 = ushort_as_bfloat16
+
+__all__ = [
+    "typing_registry",
+    "target_registry",
+    "bfloat16",
+    # Arithmetic intrinsics
+    "habs",
+    "hadd",
+    "hsub",
+    "hmul",
+    "hadd_rn",
+    "hsub_rn",
+    "hmul_rn",
+    "hdiv",
+    "hadd_sat",
+    "hsub_sat",
+    "hmul_sat",
+    "hfma",
+    "hfma_sat",
+    "hneg",
+    "hfma_relu",
+    # Comparison intrinsics
+    "heq",
+    "hne",
+    "hge",
+    "hgt",
+    "hle",
+    "hlt",
+    "hmax",
+    "hmin",
+    "hmax_nan",
+    "hmin_nan",
+    "hisinf",
+    "hisnan",
+    "hequ",
+    "hneu",
+    "hgeu",
+    "hgtu",
+    "hleu",
+    "hltu",
+    # Precision conversion and data movement
+    "float32_to_bfloat16",
+    "float64_to_bfloat16",
+    "bfloat16_to_float32",
+    "float32_to_bfloat16_rn",
+    "float32_to_bfloat16_rz",
+    "float32_to_bfloat16_rd",
+    "float32_to_bfloat16_ru",
+    "bfloat16_to_int8_rz",
+    "bfloat16_to_uint8_rz",
+    "int16_to_bfloat16_rn",
+    "int16_to_bfloat16_rz",
+    "int16_to_bfloat16_rd",
+    "int16_to_bfloat16_ru",
+    "bfloat16_to_int16_rn",
+    "bfloat16_to_int16_rz",
+    "bfloat16_to_int16_rd",
+    "bfloat16_to_int16_ru",
+    "uint16_to_bfloat16_rn",
+    "uint16_to_bfloat16_rz",
+    "uint16_to_bfloat16_rd",
+    "uint16_to_bfloat16_ru",
+    "bfloat16_to_uint16_rn",
+    "bfloat16_to_uint16_rz",
+    "bfloat16_to_uint16_rd",
+    "bfloat16_to_uint16_ru",
+    "int32_to_bfloat16_rn",
+    "int32_to_bfloat16_rz",
+    "int32_to_bfloat16_rd",
+    "int32_to_bfloat16_ru",
+    "bfloat16_to_int32_rn",
+    "bfloat16_to_int32_rz",
+    "bfloat16_to_int32_rd",
+    "bfloat16_to_int32_ru",
+    "uint32_to_bfloat16_rn",
+    "uint32_to_bfloat16_rz",
+    "uint32_to_bfloat16_rd",
+    "uint32_to_bfloat16_ru",
+    "bfloat16_to_uint32_rn",
+    "bfloat16_to_uint32_rz",
+    "bfloat16_to_uint32_rd",
+    "bfloat16_to_uint32_ru",
+    "int64_to_bfloat16_rn",
+    "int64_to_bfloat16_rz",
+    "int64_to_bfloat16_rd",
+    "int64_to_bfloat16_ru",
+    "bfloat16_to_int64_rn",
+    "bfloat16_to_int64_rz",
+    "bfloat16_to_int64_rd",
+    "bfloat16_to_int64_ru",
+    "uint64_to_bfloat16_rn",
+    "uint64_to_bfloat16_rz",
+    "uint64_to_bfloat16_rd",
+    "uint64_to_bfloat16_ru",
+    "bfloat16_to_uint64_rn",
+    "bfloat16_to_uint64_rz",
+    "bfloat16_to_uint64_rd",
+    "bfloat16_to_uint64_ru",
+    "bfloat16_as_int16",
+    "bfloat16_as_uint16",
+    "int16_as_bfloat16",
+    "uint16_as_bfloat16",
+    "htrunc",
+    "hceil",
+    "hfloor",
+    "hrint",
+    "hsqrt",
+    "hrsqrt",
+    "hrcp",
+    "hlog",
+    "hlog2",
+    "hlog10",
+    "hcos",
+    "hsin",
+    "htanh",
+    "htanh_approx",
+    "hexp",
+    "hexp2",
+    "hexp10",
+]