submission 772770

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpu L4

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

# Vector sum reduction tuned for NVIDIA L4 (Ada Lovelace, sm_89)
#
# L4 specs:
#   - 58 SMs, 128 CUDA cores/SM = 7424 cores
#   - ~300 GB/s GDDR6 bandwidth (similar to T4 but newer arch)
#   - 48 MB L2 cache (surprisingly large for the bandwidth)
#   - Max 1536 threads/SM, 48 warps/SM
#   - 72W TDP — efficiency-focused GPU
#
# Key tuning decisions:
#   - Bandwidth is the bottleneck (~300 GB/s), much lower than A100/H100
#   - 232 blocks (58 SMs × 4) for full occupancy
#   - 128 threads/block to keep register pressure low on Ada
#   - Large L2 helps with partial sum phase

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, 128>>>(
            input.data_ptr<float>(), output.data_ptr<float>(), n);
        return output;
    }

    // L4: 58 SMs, 128 threads/block
    const int threads = 128;
    int blocks;

    if (n <= 16384) {
        blocks = min(58, (n / 4 + threads - 1) / threads);
    } else {
        // 58 SMs × 4 blocks/SM = 232 blocks
        blocks = min(232, (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);

    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_l4",
    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]