numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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