Function reduce_blocks_by_key

src/backend/cuda/kernel/reduce_by_key.hpp:174–389

// Reduces each block by key
template<typename Ti, typename Tk, typename To, af_op_t op, uint DIMX>
__global__ static void reduce_blocks_by_key(int *reduced_block_sizes,
                                            Param<Tk> reduced_keys,
                                            Param<To> reduced_vals,
                                            CParam<Tk> keys, CParam<Ti> vals,
                                            int n, bool change_nan, To nanval,
                                            const int nBlocksZ) {
    const int tidx = blockIdx.x * blockDim.x + threadIdx.x;
    const int bidy = blockIdx.y;
    const int bidz = blockIdx.z % nBlocksZ;
    const int bidw = blockIdx.z / nBlocksZ;

    const int laneid = tidx % 32;

    const int nWarps = DIMX / 32;

    //
    // Allocate and initialize shared memory

    __shared__ int
        warpReduceSizes[nWarps];  // number of reduced elements in each warp

    __shared__ compute_t<Tk> warpReduceKeys[nWarps]
                                           [maxResPerWarp];  // reduced key
                                                             // segments for
                                                             // each warp
    __shared__ compute_t<To> warpReduceVals[nWarps]
                                           [maxResPerWarp];  // reduced values
                                                             // for each warp
                                                             // corresponding to
                                                             // each key segment

    // space to hold left/right-most keys of each reduced warp to check if
    // reduction should happen across boundaries
    __shared__ compute_t<Tk> warpReduceLeftBoundaryKeys[nWarps];
    __shared__ compute_t<Tk> warpReduceRightBoundaryKeys[nWarps];

    // space to hold right-most values of each reduced warp to check if
    // reduction should happen across boundaries
    __shared__ compute_t<To> warpReduceRightBoundaryVals[nWarps];

    // space to compact and finalize all reductions within block
    __shared__ compute_t<Tk> warpReduceKeysSmemFinal[nWarps * maxResPerWarp];
    __shared__ compute_t<To> warpReduceValsSmemFinal[nWarps * maxResPerWarp];

    //
    // will hold final number of reduced elements in block
    __shared__ int reducedBlockSize;

    if (threadIdx.x == 0) { reducedBlockSize = 0; }
    if (threadIdx.x < nWarps * maxResPerWarp)
        warpReduceValsSmemFinal[threadIdx.x] = scalar<compute_t<To>>(0);
    __syncthreads();

    common::Binary<compute_t<To>, op> reduce;
    common::Transform<compute_t<Ti>, compute_t<To>, op> transform;

    // load keys and values to threads
    compute_t<Tk> k;