Compiler Explorer

Source code

__global__ void batched_thread_coarsening(
    float* __restrict__ Y,
    float const* __restrict__ X,
    size_t num_elements_per_batch
) {
    size_t const NUM_THREADS{128};
    static_assert(NUM_THREADS % 32 == 0, "NUM_THREADS must be a multiple of 32.");

size_t const block_idx{blockIdx.x};
    size_t const thread_idx{threadIdx.x};
    __shared__ float shared_data[NUM_THREADS];

// Shift the input accordingly to the batch (block) index.
    X += block_idx * num_elements_per_batch;
    
    // Compute the number of elements each thread will process.
    size_t const num_elements_per_thread{(num_elements_per_batch + NUM_THREADS - 1) / NUM_THREADS};
    
    // Initialize the sum variable.
    float sum{0.0f};
    
    for (size_t i = 0; i < num_elements_per_thread; ++i) {
        size_t const offset{thread_idx + i * NUM_THREADS};
        if (offset < num_elements_per_batch)
            sum += X[offset];
    }
    shared_data[thread_idx] = sum;

#pragma unroll
    // Note that the synchronization has been moved to the beginning of the loop below.
    for (size_t stride = NUM_THREADS / 2; stride > 0; stride >>= 1) {
        __syncthreads();
        if (thread_idx < stride)
            shared_data[thread_idx] += shared_data[thread_idx + stride];
    }

if (thread_idx == 0)
        Y[block_idx] = shared_data[0];
}

Source code

__device__ void warp_reduce(volatile float* shared_data, size_t thread_idx) {
    shared_data[thread_idx] += shared_data[thread_idx + 32];
    shared_data[thread_idx] += shared_data[thread_idx + 16];
    shared_data[thread_idx] += shared_data[thread_idx + 8];
    shared_data[thread_idx] += shared_data[thread_idx + 4];
    shared_data[thread_idx] += shared_data[thread_idx + 2];
    shared_data[thread_idx] += shared_data[thread_idx + 1];
}

__global__ void batched_fully_unroll(
    float* __restrict__ Y,
    float const* __restrict__ X,
    size_t num_elements_per_batch
) {
    size_t const NUM_THREADS{128};
    size_t const NUM_THREADS_PER_WRAP{32};
    static_assert(NUM_THREADS % 32 == 0, "NUM_THREADS must be a multiple of 32.");

size_t const block_idx{blockIdx.x};
    size_t const thread_idx{threadIdx.x};
    __shared__ float shared_data[NUM_THREADS];

// Handle elements of the indices > thread index.
    for (size_t i = 0; i < num_elements_per_thread; ++i) {
        size_t const offset{thread_idx + i * NUM_THREADS};
        if (offset < num_elements_per_batch)
            sum += X[offset];
    }
    shared_data[thread_idx] = sum;

for (size_t stride = NUM_THREADS / 2; stride > NUM_THREADS_PER_WRAP; stride >>= 1) {
        __syncthreads();
        if (thread_idx < stride)
            shared_data[thread_idx] += shared_data[thread_idx + stride];
    }

if (thread_idx < NUM_THREADS_PER_WRAP) {
        __syncthreads();
        warp_reduce(shared_data, thread_idx);
    }

if (thread_idx == 0)
        Y[block_idx] = shared_data[0];
}

Source code

__device__ void warp_reduce(volatile float* shared_data, size_t thread_idx, size_t const NUM_THREADS) {
    if (NUM_THREADS >= 64) shared_data[thread_idx] += shared_data[thread_idx + 32];
    if (NUM_THREADS >= 32) shared_data[thread_idx] += shared_data[thread_idx + 16];
    if (NUM_THREADS >= 16) shared_data[thread_idx] += shared_data[thread_idx + 8];
    if (NUM_THREADS >= 8) shared_data[thread_idx] += shared_data[thread_idx + 4];
    if (NUM_THREADS >= 4) shared_data[thread_idx] += shared_data[thread_idx + 2];
    if (NUM_THREADS >= 2) shared_data[thread_idx] += shared_data[thread_idx + 1];
}

size_t const block_idx{blockIdx.x};
    size_t const thread_idx{threadIdx.x};
    __shared__ float shared_data[NUM_THREADS];

if (NUM_THREADS == 1024){
        if (thread_idx < 512) shared_data[thread_idx] += shared_data[thread_idx + 512];
        __syncthreads();
    }
    if (NUM_THREADS >= 512){
        if (thread_idx < 256) shared_data[thread_idx] += shared_data[thread_idx + 256];
        __syncthreads();
    }
    if (NUM_THREADS >= 256){
        if (thread_idx < 128) shared_data[thread_idx] += shared_data[thread_idx + 128];
        __syncthreads();
    }
    if (NUM_THREADS >= 128){
        if (thread_idx < 64) shared_data[thread_idx] += shared_data[thread_idx + 64];
        __syncthreads();
    }

if (thread_idx < NUM_THREADS_PER_WRAP)
        warp_reduce(shared_data, thread_idx, NUM_THREADS);

if (thread_idx == 0)
        Y[block_idx] = shared_data[0];
}