submission 772774

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpu B200

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

# Vector sum reduction tuned for NVIDIA B200 (Blackwell, sm_100)
#
# B200 specs:
#   - 192 SMs, 128 CUDA cores/SM = 24576 cores
#   - ~8 TB/s HBM3e bandwidth (monster)
#   - 96 MB L2 cache
#   - Max 2048 threads/SM, 64 warps/SM
#
# Key tuning decisions:
#   - Massive bandwidth means kernel launch overhead is relatively more significant
#   - 768 blocks (192 SMs × 4) for full occupancy
#   - 256 threads/block for high throughput
#   - With 8 TB/s bandwidth, even 52M floats (~200MB) finishes in ~25µs theoretically
#   - Must minimize any overhead: launch, allocation, synchronization

cuda_source = r"""
#include <cuda_runtime.h>

__device__ __forceinline__ float warp_reduce_sum(float val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void reduce_tiny(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n
) {
    __shared__ float smem[32];
    const int tid = threadIdx.x;
    float sum = 0.0f;

    for (int i = tid; i < n; i += blockDim.x) {
        sum += input[i];
    }

    sum = warp_reduce_sum(sum);
    const int lane = tid & 31;
    const int warp_id = tid >> 5;

    if (lane == 0) smem[warp_id] = sum;
    __syncthreads();

    const int num_warps = (blockDim.x + 31) / 32;
    if (warp_id == 0) {
        sum = (lane < num_warps) ? smem[lane] : 0.0f;
        sum = warp_reduce_sum(sum);
    }
    if (tid == 0) output[0] = sum;
}

__global__ void reduce_phase1(
    const float* __restrict__ input,
    float* __restrict__ partial_sums,
    int n
) {
    __shared__ float smem[32];
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int block_size = blockDim.x;
    const int grid_size = gridDim.x * block_size;

    float sum = 0.0f;

    const int n4 = n >> 2;
    const float4* input4 = reinterpret_cast<const float4*>(input);
    for (int i = bid * block_size + tid; i < n4; i += grid_size) {
        float4 v = input4[i];
        sum += v.x + v.y + v.z + v.w;
    }

    int tail_start = n4 << 2;
    for (int i = tail_start + bid * block_size + tid; i < n; i += grid_size) {
        sum += input[i];
    }

    sum = warp_reduce_sum(sum);
    const int lane = tid & 31;
    const int warp_id = tid >> 5;

    if (lane == 0) smem[warp_id] = sum;
    __syncthreads();

    const int num_warps = (block_size + 31) / 32;
    if (warp_id == 0) {
        sum = (lane < num_warps) ? smem[lane] : 0.0f;
        sum = warp_reduce_sum(sum);
    }
    if (tid == 0) partial_sums[bid] = sum;
}

__global__ void reduce_phase2(
    const float* __restrict__ partial_sums,
    float* __restrict__ output,
    int n
) {
    __shared__ float smem[32];
    const int tid = threadIdx.x;
    float sum = 0.0f;

    for (int i = tid; i < n; i += blockDim.x) {
        sum += partial_sums[i];
    }

    sum = warp_reduce_sum(sum);
    const int lane = tid & 31;
    const int warp_id = tid >> 5;

    if (lane == 0) smem[warp_id] = sum;
    __syncthreads();

    const int num_warps = (blockDim.x + 31) / 32;
    if (warp_id == 0) {
        sum = (lane < num_warps) ? smem[lane] : 0.0f;
        sum = warp_reduce_sum(sum);
    }
    if (tid == 0) output[0] = sum;
}

torch::Tensor vector_sum_cuda(torch::Tensor input, torch::Tensor output) {
    const int n = input.numel();

    if (n <= 1024) {
        reduce_tiny<<<1, 256>>>(
            input.data_ptr<float>(), output.data_ptr<float>(), n);
        return output;
    }

    // B200: 192 SMs, 256 threads/block
    const int threads = 256;
    int blocks;

    if (n <= 32768) {
        blocks = min(192, (n / 4 + threads - 1) / threads);
    } else {
        // 192 SMs × 4 blocks/SM = 768 blocks
        blocks = min(768, (n / 4 + threads - 1) / threads);
    }

    auto partial_sums = torch::empty({blocks}, input.options());

    reduce_phase1<<<blocks, threads>>>(
        input.data_ptr<float>(), partial_sums.data_ptr<float>(), n);

    // 768 partial sums → one block of 256 threads handles it in 3 iterations
    int p2_threads = min(256, ((blocks + 31) / 32) * 32);
    reduce_phase2<<<1, p2_threads>>>(
        partial_sums.data_ptr<float>(), output.data_ptr<float>(), blocks);

    return output;
}
"""

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor vector_sum_cuda(torch::Tensor input, torch::Tensor output);
"""

module = load_inline(
    name="vector_sum_b200",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["vector_sum_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


def custom_kernel(data: input_t) -> output_t:
    data, output = data
    module.vector_sum_cuda(data, output)
    return output[0]