[FEA]: Allow sub-partition of a block reduction

[nv-cuchytil]

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this request and that I agree to the Code of Conduct

Area

CUB

Is your feature request related to a problem? Please describe.

Block reductions return a single value in the first thread, but there are instances were N results are needed in different threads. At a warp level, cub allows you to create "logical" warps, which partition a hardware thread group further, but this functionality is missing at a block level. All non-exited threads have to participate in block API call.

Describe the solution you'd like

Two possible solutions:

• Provide segmented / batched block reduction
• Allow sub-partitions of a block

Describe alternatives you've considered

The current work-around is to utilize cub's warp reduction and implement something custom to then combine the warp level reductions for the block but it would be nice for this to be a single cub API call.

Additional context

No response

[github-actions]

Hi @nv-cuchytil!

Thanks for submitting this issue - the CCCL team has been notified and we'll get back to you as soon as we can!
In the mean time, feel free to add any relevant information to this issue.

[nv-cuchytil]

Here's somewhat of a contrived example illustrating the request:

#include <cstdio>
#include <vector>
#include <iostream>
#include <iomanip>
#include <cuda_runtime.h>

#include <cub/block/block_reduce.cuh>
#include <cub/warp/warp_reduce.cuh>
#include <cub/util_device.cuh>
#include <cstdint>

constexpr unsigned int THREADS_PER_BLOCK = 128;
constexpr unsigned int WARP_SIZE = 32;

constexpr unsigned int SIZE = 1024;
constexpr unsigned int GROUPS_PER_BLOCK = 2;

constexpr unsigned int ELEMENTS = SIZE / (THREADS_PER_BLOCK / GROUPS_PER_BLOCK);
constexpr unsigned int NUM_WARPS = THREADS_PER_BLOCK / WARP_SIZE;
constexpr unsigned int WARPS_PER_GROUP = NUM_WARPS / GROUPS_PER_BLOCK;
// constexpr unsigned int THREADS_PER_GROUP = THREADS_PER_BLOCK / GROUPS_PER_BLOCK;

constexpr unsigned int BATCH_SIZE = 2;

static_assert(NUM_WARPS >= GROUPS_PER_BLOCK);
// ensure GROUPS_PER_BLOCK is power of 2
static_assert((GROUPS_PER_BLOCK & (GROUPS_PER_BLOCK - 1)) == 0);

#define CUDA_CHECK(call) do { \
    cudaError_t err = (call); \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA error at %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
        exit(EXIT_FAILURE); \
    } \
} while(0)

using ValueT = float;
using WarpReduceT = cub::WarpReduce<ValueT>;

__device__ void cub_block_reduce_sum(ValueT* input_rmem, ValueT* output_rmem) {
    __shared__ ValueT shared_sum[NUM_WARPS];

    ValueT thread_data[ELEMENTS];
    #pragma unroll
    for (unsigned int i = 0; i < ELEMENTS; ++i) {
        thread_data[i] = input_rmem[i];
    }

    // if constexpr (WARPS_PER_GROUP == NUM_WARPS) {
    //     using BlockReduceT = cub::BlockReduce<ValueT, NUM_THREADS>;
    //     __shared__ typename BlockReduceT::TempStorage temp_storage;

    //     ValueT sum = BlockReduceT(temp_storage).Sum(thread_data);
    //     if (threadIdx.x == 0) {
    //         shared_sum[0] = sum;
    //     }
    //     __syncthreads();
    // }
    // else {
        __shared__ typename WarpReduceT::TempStorage temp_storage[NUM_WARPS];

        const unsigned int lane = threadIdx.x % WARP_SIZE;
        const unsigned int warp = threadIdx.x / WARP_SIZE;

        const unsigned int group_lane = warp % WARPS_PER_GROUP;
        const unsigned int group_idx = warp / GROUPS_PER_BLOCK;

        ValueT warp_sum = WarpReduceT(temp_storage[warp]).Sum(thread_data);

        // first thread of each warp writes its warp sum back into shared memory
        if (lane == 0) {
            shared_sum[warp] = warp_sum;
        }

        __syncthreads();

        // if there is more than 1 warp per FFT group, compute inter-FFT group reduction
        if constexpr (WARPS_PER_GROUP > 1) {
            // first N threads of first warp of each fft_group load values from shared and perform a warp reduce to get final reduced values
            thread_data[0] = (lane < WARPS_PER_GROUP && group_lane == 0) ? shared_sum[lane + group_idx * GROUPS_PER_BLOCK] : static_cast<ValueT>(0);

            warp_sum = WarpReduceT(temp_storage[warp]).Sum(thread_data[0]);

            if (lane == 0 && group_lane == 0) {
                shared_sum[warp] = warp_sum;
            }

            __syncthreads();
        }
    // }

    #pragma unroll
    for (unsigned int n = 0; n < GROUPS_PER_BLOCK; ++n) {
        output_rmem[n] = shared_sum[n * WARPS_PER_GROUP];
    }
}

__global__ void test_kernel(float* in, float* in_energy) {
    float  energy_rmem_in[ELEMENTS];
    float  energy_rmem_out[GROUPS_PER_BLOCK];

    float* batch_in         = in         + blockIdx.x * GROUPS_PER_BLOCK * SIZE;
    float* batch_in_energy  = in_energy  + blockIdx.x * GROUPS_PER_BLOCK;

    const unsigned int group_lane = threadIdx.x % (WARP_SIZE * GROUPS_PER_BLOCK);
    const unsigned int group_idx = threadIdx.x / (WARP_SIZE * GROUPS_PER_BLOCK);

    // load input into groups of warps
    #pragma unroll
    for (unsigned int i = 0; i < ELEMENTS; ++i) {
        energy_rmem_in[i] = batch_in[group_lane + (i * WARP_SIZE * GROUPS_PER_BLOCK) + (group_idx * SIZE)];
    }

    cub_block_reduce_sum(energy_rmem_in, energy_rmem_out);
    if (threadIdx.x == 0) {
        for (unsigned int n = 0; n < GROUPS_PER_BLOCK; ++n) {
            batch_in_energy[n] = energy_rmem_out[n];
        }
    }
}

int main() {
    constexpr size_t total_input_elements = SIZE * BATCH_SIZE;
    constexpr size_t total_input_bytes = total_input_elements * sizeof(float);

    constexpr size_t total_output_elements = BATCH_SIZE;
    constexpr size_t total_output_bytes = total_output_elements * sizeof(float);

    std::vector<float> h_input(total_input_elements);
    std::vector<float> h_output(total_output_elements);

    float answers[BATCH_SIZE] = { 0.0 };

    // Initialize with random data
    for (size_t i = 0; i < total_input_elements; ++i) {
        h_input[i] = i;
        answers[i / SIZE] += i;
    }

    float* d_input;
    float* d_output;
    CUDA_CHECK(cudaMalloc(&d_input, total_input_bytes));
    CUDA_CHECK(cudaMalloc(&d_output, total_output_bytes));
    
    CUDA_CHECK(cudaMemcpy(d_input, h_input.data(), total_input_bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemset(d_output, 0, total_output_bytes));

    dim3 grid(BATCH_SIZE / GROUPS_PER_BLOCK);
    dim3 block(THREADS_PER_BLOCK);

    test_kernel<<<grid, block>>>(d_input, d_output);
    CUDA_CHECK(cudaDeviceSynchronize());
    
    CUDA_CHECK(cudaMemcpy(h_output.data(), d_output, total_output_bytes, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_input));
    CUDA_CHECK(cudaFree(d_output));

    std::cout << std::string(70, '=') << "\n";

    for (unsigned int i = 0; i < BATCH_SIZE; ++i) {
        float delta = abs(h_output[i] - answers[i]);
        std::cout << h_output[i] << ((delta < 1e-6) ? " == " : " != ") << answers[i] << "\n";
    }

    std::cout << std::string(70, '=') << "\n";

    return 0;
}