submission 780494

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpu B200

# Sum reduction. Reference returns a 0-dim scalar tensor (shape ()), so we
# must return a 0-dim view of the (1,) output tensor to match shape.
# Strategy: warp-shuffle reduce → block reduce via shared mem → atomicAdd
# to scalar output. Vectorized float4 loads for the bulk + scalar tail.
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t


_CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <cstdint>

__global__ void sum_reduce_f32(const float* __restrict__ x,
                               float* __restrict__ out_scalar,
                               int n) {
    __shared__ float smem[32];
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    float local = 0.0f;
    int n4 = n & ~3;
    for (int i = tid * 4; i < n4; i += stride * 4) {
        float4 v = *reinterpret_cast<const float4*>(x + i);
        local += v.x + v.y + v.z + v.w;
    }
    for (int i = n4 + tid; i < n; i += stride) local += x[i];

    for (int off = 16; off > 0; off >>= 1)
        local += __shfl_xor_sync(0xffffffff, local, off);
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    if (lane == 0) smem[warp] = local;
    __syncthreads();
    if (warp == 0) {
        int num_warps = (blockDim.x + 31) >> 5;
        local = (lane < num_warps) ? smem[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            local += __shfl_xor_sync(0xffffffff, local, off);
        if (lane == 0) atomicAdd(out_scalar, local);
    }
}

void launch_sum(uintptr_t x_ptr, uintptr_t scalar_ptr, int n) {
    int threads = 256;
    int blocks = (n + threads * 4 - 1) / (threads * 4);
    if (blocks > 1024) blocks = 1024;
    if (blocks < 1) blocks = 1;
    sum_reduce_f32<<<blocks, threads>>>(
        reinterpret_cast<const float*>(x_ptr),
        reinterpret_cast<float*>(scalar_ptr), n);
}
"""

_CPP_SRC = """
void launch_sum(uintptr_t, uintptr_t, int);
"""

_mod = load_inline(
    name="vectorsum_v5",
    cpp_sources=_CPP_SRC, cuda_sources=_CUDA_SRC,
    functions=["launch_sum"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_100"],
    extra_cflags=["-O3"], verbose=False)


def custom_kernel(data: input_t) -> output_t:
    x, out = data[0], data[-1]
    if not x.is_contiguous():
        x = x.contiguous()
    n = x.numel()

    if x.dtype == torch.float32 and out.dtype == torch.float32:
        out.zero_()
        _mod.launch_sum(x.data_ptr(), out.data_ptr(), n)
    else:
        s = x.to(torch.float32).sum()
        out.view(-1)[0] = s.to(out.dtype)

    return out.view(())