reduction.py

"""
A library written in CUDA Python for generating reduction kernels
"""

from numba.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]