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

numba_cuda/numba/cuda/itanium_mangler.py

@@ -0,0 +1,214 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Itanium CXX ABI Mangler
+
+Reference: https://itanium-cxx-abi.github.io/cxx-abi/abi.html
+
+The basics of the mangling scheme.
+
+We are hijacking the CXX mangling scheme for our use.  We map Python modules
+into CXX namespace.  A `module1.submodule2.foo` is mapped to
+`module1::submodule2::foo`.   For parameterized numba types, we treat them as
+templated types; for example, `array(int64, 1d, C)` becomes an
+`array<int64, 1, C>`.
+
+All mangled names are prefixed with "_Z".  It is followed by the name of the
+entity.  A name contains one or more identifiers.  Each identifier is encoded
+as "<num of char><name>".   If the name is namespaced and, therefore,
+has multiple identifiers, the entire name is encoded as "N<name>E".
+
+For functions, arguments types follow.  There are condensed encodings for basic
+built-in types; e.g. "i" for int, "f" for float.  For other types, the
+previously mentioned name encoding should be used.
+
+For templated types, the template parameters are encoded immediately after the
+name.  If it is namespaced, it should be within the 'N' 'E' marker.  Template
+parameters are encoded in "I<params>E", where each parameter is encoded using
+the mentioned name encoding scheme.  Template parameters can contain literal
+values like the '1' in the array type shown earlier.  There is special encoding
+scheme for them to avoid leading digits.
+"""
+
+import re
+
+from numba.cuda import types
+
+
+# According the scheme, valid characters for mangled names are [a-zA-Z0-9_].
+# We borrow the '_' as the escape character to encode invalid char into
+# '_xx' where 'xx' is the hex codepoint.
+_re_invalid_char = re.compile(r"[^a-z0-9_]", re.I)
+
+PREFIX = "_Z"
+
+# Numba types to mangled type code. These correspond with the codes listed in
+# https://itanium-cxx-abi.github.io/cxx-abi/abi.html#mangling-builtin
+N2CODE = {
+    types.void: "v",
+    types.boolean: "b",
+    types.uint8: "h",
+    types.int8: "a",
+    types.uint16: "t",
+    types.int16: "s",
+    types.uint32: "j",
+    types.int32: "i",
+    types.uint64: "y",
+    types.int64: "x",
+    types.float16: "Dh",
+    types.float32: "f",
+    types.float64: "d",
+}
+
+
+def _escape_string(text):
+    """Escape the given string so that it only contains ASCII characters
+    of [a-zA-Z0-9_$].
+
+    The dollar symbol ($) and other invalid characters are escaped into
+    the string sequence of "$xx" where "xx" is the hex codepoint of the char.
+
+    Multibyte characters are encoded into utf8 and converted into the above
+    hex format.
+    """
+
+    def repl(m):
+        return "".join(("_%02x" % ch) for ch in m.group(0).encode("utf8"))
+
+    ret = re.sub(_re_invalid_char, repl, text)
+    # Return str if we got a unicode (for py2)
+    if not isinstance(ret, str):
+        return ret.encode("ascii")
+    return ret
+
+
+def _fix_lead_digit(text):
+    """
+    Fix text with leading digit
+    """
+    if text and text[0].isdigit():
+        return "_" + text
+    else:
+        return text
+
+
+def _len_encoded(string):
+    """
+    Prefix string with digit indicating the length.
+    Add underscore if string is prefixed with digits.
+    """
+    string = _fix_lead_digit(string)
+    return "%u%s" % (len(string), string)
+
+
+def mangle_abi_tag(abi_tag: str) -> str:
+    return "B" + _len_encoded(_escape_string(abi_tag))
+
+
+def mangle_identifier(ident, template_params="", *, abi_tags=(), uid=None):
+    """
+    Mangle the identifier with optional template parameters and abi_tags.
+
+    Note:
+
+    This treats '.' as '::' in C++.
+    """
+    if uid is not None:
+        # Add uid to abi-tags
+        abi_tags = (f"v{uid}", *abi_tags)
+    parts = [_len_encoded(_escape_string(x)) for x in ident.split(".")]
+    enc_abi_tags = list(map(mangle_abi_tag, abi_tags))
+    extras = template_params + "".join(enc_abi_tags)
+    if len(parts) > 1:
+        return "N%s%sE" % ("".join(parts), extras)
+    else:
+        return "%s%s" % (parts[0], extras)
+
+
+def mangle_type_or_value(typ):
+    """
+    Mangle type parameter and arbitrary value.
+    """
+    # Handle numba types
+    if isinstance(typ, types.Type):
+        if typ in N2CODE:
+            return N2CODE[typ]
+        else:
+            return mangle_templated_ident(*typ.mangling_args)
+    # Handle integer literal
+    elif isinstance(typ, int):
+        return "Li%dE" % typ
+    # Handle str as identifier
+    elif isinstance(typ, str):
+        return mangle_identifier(typ)
+    # Otherwise
+    else:
+        enc = _escape_string(str(typ))
+        return _len_encoded(enc)
+
+
+# Alias
+mangle_type = mangle_type_or_value
+mangle_value = mangle_type_or_value
+
+
+def mangle_templated_ident(identifier, parameters):
+    """
+    Mangle templated identifier.
+    """
+    template_params = (
+        "I%sE" % "".join(map(mangle_type_or_value, parameters))
+        if parameters
+        else ""
+    )
+    return mangle_identifier(identifier, template_params)
+
+
+def mangle_args(argtys):
+    """
+    Mangle sequence of Numba type objects and arbitrary values.
+    """
+    return "".join([mangle_type_or_value(t) for t in argtys])
+
+
+def mangle(ident, argtys, *, abi_tags=(), uid=None):
+    """
+    Mangle identifier with Numba type objects and abi-tags.
+    """
+    return "".join(
+        [
+            PREFIX,
+            mangle_identifier(ident, abi_tags=abi_tags, uid=uid),
+            mangle_args(argtys),
+        ]
+    )
+
+
+def prepend_namespace(mangled, ns):
+    """
+    Prepend namespace to mangled name.
+    """
+    if not mangled.startswith(PREFIX):
+        raise ValueError("input is not a mangled name")
+    elif mangled.startswith(PREFIX + "N"):
+        # nested
+        remaining = mangled[3:]
+        ret = PREFIX + "N" + mangle_identifier(ns) + remaining
+    else:
+        # non-nested
+        remaining = mangled[2:]
+        head, tail = _split_mangled_ident(remaining)
+        ret = PREFIX + "N" + mangle_identifier(ns) + head + "E" + tail
+    return ret
+
+
+def _split_mangled_ident(mangled):
+    """
+    Returns `(head, tail)` where `head` is the `<len> + <name>` encoded
+    identifier and `tail` is the remaining.
+    """
+    ct = int(mangled)
+    ctlen = len(str(ct))
+    at = ctlen + ct
+    return mangled[:at], mangled[at:]

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

@@ -0,0 +1,265 @@
+# 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
+"""
+
+from numba.cuda.np.numpy_support import from_dtype
+
+
+_WARPSIZE = 32
+_NUMWARPS = 4
+
+
+def _gpu_reduce_factory(fn, nbtype):
+    from numba import cuda
+
+    reduce_op = cuda.jit(device=True)(fn)
+    inner_sm_size = _WARPSIZE + 1  # plus one to avoid SM collision
+    max_blocksize = _NUMWARPS * _WARPSIZE
+
+    @cuda.jit(device=True)
+    def inner_warp_reduction(sm_partials, init):
+        """
+        Compute reduction within a single warp
+        """
+        tid = cuda.threadIdx.x
+        warpid = tid // _WARPSIZE
+        laneid = tid % _WARPSIZE
+
+        sm_this = sm_partials[warpid, :]
+        sm_this[laneid] = init
+        cuda.syncwarp()
+
+        width = _WARPSIZE // 2
+        while width:
+            if laneid < width:
+                old = sm_this[laneid]
+                sm_this[laneid] = reduce_op(old, sm_this[laneid + width])
+            cuda.syncwarp()
+            width //= 2
+
+    @cuda.jit(device=True)
+    def device_reduce_full_block(arr, partials, sm_partials):
+        """
+        Partially reduce `arr` into `partials` using `sm_partials` as working
+        space.  The algorithm goes like:
+
+            array chunks of 128:  |   0 | 128 | 256 | 384 | 512 |
+                        block-0:  |   x |     |     |   x |     |
+                        block-1:  |     |   x |     |     |   x |
+                        block-2:  |     |     |   x |     |     |
+
+        The array is divided into chunks of 128 (size of a threadblock).
+        The threadblocks consumes the chunks in roundrobin scheduling.
+        First, a threadblock loads a chunk into temp memory.  Then, all
+        subsequent chunks are combined into the temp memory.
+
+        Once all chunks are processed.  Inner-block reduction is performed
+        on the temp memory.  So that, there will just be one scalar result
+        per block.  The result from each block is stored to `partials` at
+        the dedicated slot.
+        """
+        tid = cuda.threadIdx.x
+        blkid = cuda.blockIdx.x
+        blksz = cuda.blockDim.x
+        gridsz = cuda.gridDim.x
+
+        # block strided loop to compute the reduction
+        start = tid + blksz * blkid
+        stop = arr.size
+        step = blksz * gridsz
+
+        # load first value
+        tmp = arr[start]
+        # loop over all values in block-stride
+        for i in range(start + step, stop, step):
+            tmp = reduce_op(tmp, arr[i])
+
+        cuda.syncthreads()
+        # inner-warp reduction
+        inner_warp_reduction(sm_partials, tmp)
+
+        cuda.syncthreads()
+        # at this point, only the first slot for each warp in tsm_partials
+        # is valid.
+
+        # finish up block reduction
+        # warning: this is assuming 4 warps.
+        # assert numwarps == 4
+        if tid < 2:
+            sm_partials[tid, 0] = reduce_op(
+                sm_partials[tid, 0], sm_partials[tid + 2, 0]
+            )
+            cuda.syncwarp()
+        if tid == 0:
+            partials[blkid] = reduce_op(sm_partials[0, 0], sm_partials[1, 0])
+
+    @cuda.jit(device=True)
+    def device_reduce_partial_block(arr, partials, sm_partials):
+        """
+        This computes reduction on `arr`.
+        This device function must be used by 1 threadblock only.
+        The blocksize must match `arr.size` and must not be greater than 128.
+        """
+        tid = cuda.threadIdx.x
+        blkid = cuda.blockIdx.x
+        blksz = cuda.blockDim.x
+        warpid = tid // _WARPSIZE
+        laneid = tid % _WARPSIZE
+
+        size = arr.size
+        # load first value
+        tid = cuda.threadIdx.x
+        value = arr[tid]
+        sm_partials[warpid, laneid] = value
+
+        cuda.syncthreads()
+
+        if (warpid + 1) * _WARPSIZE < size:
+            # fully populated warps
+            inner_warp_reduction(sm_partials, value)
+        else:
+            # partially populated warps
+            # NOTE: this uses a very inefficient sequential algorithm
+            if laneid == 0:
+                sm_this = sm_partials[warpid, :]
+                base = warpid * _WARPSIZE
+                for i in range(1, size - base):
+                    sm_this[0] = reduce_op(sm_this[0], sm_this[i])
+
+        cuda.syncthreads()
+        # finish up
+        if tid == 0:
+            num_active_warps = (blksz + _WARPSIZE - 1) // _WARPSIZE
+
+            result = sm_partials[0, 0]
+            for i in range(1, num_active_warps):
+                result = reduce_op(result, sm_partials[i, 0])
+
+            partials[blkid] = result
+
+    def gpu_reduce_block_strided(arr, partials, init, use_init):
+        """
+        Perform reductions on *arr* and writing out partial reduction result
+        into *partials*.  The length of *partials* is determined by the
+        number of threadblocks. The initial value is set with *init*.
+
+        Launch config:
+
+        Blocksize must be multiple of warpsize and it is limited to 4 warps.
+        """
+        tid = cuda.threadIdx.x
+
+        sm_partials = cuda.shared.array(
+            (_NUMWARPS, inner_sm_size), dtype=nbtype
+        )
+        if cuda.blockDim.x == max_blocksize:
+            device_reduce_full_block(arr, partials, sm_partials)
+        else:
+            device_reduce_partial_block(arr, partials, sm_partials)
+        # deal with the initializer
+        if use_init and tid == 0 and cuda.blockIdx.x == 0:
+            partials[0] = reduce_op(partials[0], init)
+
+    return cuda.jit(gpu_reduce_block_strided)
+
+
+class Reduce(object):
+    """Create a reduction object that reduces values using a given binary
+    function. The binary function is compiled once and cached inside this
+    object. Keeping this object alive will prevent re-compilation.
+    """
+
+    _cache = {}
+
+    def __init__(self, functor):
+        """
+        :param functor: A function implementing a binary operation for
+                        reduction. It will be compiled as a CUDA device
+                        function using ``cuda.jit(device=True)``.
+        """
+        self._functor = functor
+
+    def _compile(self, dtype):
+        key = self._functor, dtype
+        if key in self._cache:
+            kernel = self._cache[key]
+        else:
+            kernel = _gpu_reduce_factory(self._functor, from_dtype(dtype))
+            self._cache[key] = kernel
+        return kernel
+
+    def __call__(self, arr, size=None, res=None, init=0, stream=0):
+        """Performs a full reduction.
+
+        :param arr: A host or device array.
+        :param size: Optional integer specifying the number of elements in
+                    ``arr`` to reduce. If this parameter is not specified, the
+                    entire array is reduced.
+        :param res: Optional device array into which to write the reduction
+                    result to. The result is written into the first element of
+                    this array. If this parameter is specified, then no
+                    communication of the reduction output takes place from the
+                    device to the host.
+        :param init: Optional initial value for the reduction, the type of which
+                    must match ``arr.dtype``.
+        :param stream: Optional CUDA stream in which to perform the reduction.
+                    If no stream is specified, the default stream of 0 is
+                    used.
+        :return: If ``res`` is specified, ``None`` is returned. Otherwise, the
+                result of the reduction is returned.
+        """
+        from numba import cuda
+
+        # ensure 1d array
+        if arr.ndim != 1:
+            raise TypeError("only support 1D array")
+
+        # adjust array size
+        if size is not None:
+            arr = arr[:size]
+
+        init = arr.dtype.type(init)  # ensure the right type
+
+        # return `init` if `arr` is empty
+        if arr.size < 1:
+            return init
+
+        kernel = self._compile(arr.dtype)
+
+        # Perform the reduction on the GPU
+        blocksize = _NUMWARPS * _WARPSIZE
+        size_full = (arr.size // blocksize) * blocksize
+        size_partial = arr.size - size_full
+        full_blockct = min(size_full // blocksize, _WARPSIZE * 2)
+
+        # allocate size of partials array
+        partials_size = full_blockct
+        if size_partial:
+            partials_size += 1
+        partials = cuda.device_array(shape=partials_size, dtype=arr.dtype)
+
+        if size_full:
+            # kernel for the fully populated threadblocks
+            kernel[full_blockct, blocksize, stream](
+                arr[:size_full], partials[:full_blockct], init, True
+            )
+
+        if size_partial:
+            # kernel for partially populated threadblocks
+            kernel[1, size_partial, stream](
+                arr[size_full:], partials[full_blockct:], init, not full_blockct
+            )
+
+        if partials.size > 1:
+            # finish up
+            kernel[1, partials_size, stream](partials, partials, init, False)
+
+        # handle return value
+        if res is not None:
+            res[:1].copy_to_device(partials[:1], stream=stream)
+            return
+        else:
+            return partials[0]

numba_cuda/numba/cuda/kernels/transpose.py

@@ -0,0 +1,65 @@
+# 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
+from numba.cuda.np import numpy_support as nps
+
+
+def transpose(a, b=None):
+    """Compute the transpose of 'a' and store it into 'b', if given,
+    and return it. If 'b' is not given, allocate a new array
+    and return that.
+
+    This implements the algorithm documented in
+    http://devblogs.nvidia.com/parallelforall/efficient-matrix-transpose-cuda-cc/
+
+    :param a: an `np.ndarray` or a `DeviceNDArrayBase` subclass. If already on
+        the device its stream will be used to perform the transpose (and to copy
+        `b` to the device if necessary).
+    """
+
+    # prefer `a`'s stream if
+    stream = getattr(a, "stream", 0)
+
+    if not b:
+        cols, rows = a.shape
+        strides = a.dtype.itemsize * cols, a.dtype.itemsize
+        b = cuda.cudadrv.devicearray.DeviceNDArray(
+            (rows, cols), strides, dtype=a.dtype, stream=stream
+        )
+
+    dt = nps.from_dtype(a.dtype)
+
+    tpb = driver.get_device().MAX_THREADS_PER_BLOCK
+    # we need to factor available threads into x and y axis
+    tile_width = int(math.pow(2, math.log(tpb, 2) / 2))
+    tile_height = int(tpb / tile_width)
+
+    tile_shape = (tile_height, tile_width + 1)
+
+    @cuda.jit
+    def kernel(input, output):
+        tile = cuda.shared.array(shape=tile_shape, dtype=dt)
+
+        tx = cuda.threadIdx.x
+        ty = cuda.threadIdx.y
+        bx = cuda.blockIdx.x * cuda.blockDim.x
+        by = cuda.blockIdx.y * cuda.blockDim.y
+        x = by + tx
+        y = bx + ty
+
+        if by + ty < input.shape[0] and bx + tx < input.shape[1]:
+            tile[ty, tx] = input[by + ty, bx + tx]
+        cuda.syncthreads()
+        if y < output.shape[0] and x < output.shape[1]:
+            output[y, x] = tile[tx, ty]
+
+    # one block per tile, plus one for remainders
+    blocks = int(b.shape[0] / tile_height + 1), int(b.shape[1] / tile_width + 1)
+    # one thread per tile element
+    threads = tile_height, tile_width
+    kernel[blocks, threads, stream](a, b)
+
+    return b