submission 780504

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

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

// first pass: each block reduces BLOCK_SIZE fp32 elements into one fp64
// partial sum. casting to double inside the load is what gives us the
// precision the reference implementation gets from .to(float64).sum().
template <int BLOCK_SIZE>
__global__ void reduce_kernel(
    const float* __restrict__ data,
    double* __restrict__ partial,
    int n
) {
    __shared__ double smem[BLOCK_SIZE];

    int tid = threadIdx.x;
    int gid = blockIdx.x * BLOCK_SIZE + tid;

    smem[tid] = (gid < n) ? static_cast<double>(data[gid]) : 0.0;
    __syncthreads();

    // tree reduction in shared memory, all in fp64
    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (tid < stride)
            smem[tid] += smem[tid + stride];
        __syncthreads();
    }

    if (tid == 0)
        partial[blockIdx.x] = smem[0];
}

// second pass: a single block sums the fp64 partials and writes the
// final fp32 result. the loop lets one block handle more than
// BLOCK_SIZE partials when n is large.
template <int BLOCK_SIZE>
__global__ void final_reduce_kernel(
    const double* __restrict__ partial,
    float* __restrict__ output,
    int n
) {
    __shared__ double smem[BLOCK_SIZE];

    int tid = threadIdx.x;

    double val = 0.0;
    for (int i = tid; i < n; i += BLOCK_SIZE)
        val += partial[i];
    smem[tid] = val;
    __syncthreads();

    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (tid < stride)
            smem[tid] += smem[tid + stride];
        __syncthreads();
    }

    if (tid == 0)
        output[0] = static_cast<float>(smem[0]);
}

torch::Tensor vectorsum_cuda(torch::Tensor data, torch::Tensor output) {
    int n = data.numel();
    const int BLOCK_SIZE = 256;
    int num_blocks = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;

    // fp64 scratch buffer for per-block partial sums
    auto partial = torch::empty(
        {num_blocks},
        data.options().dtype(torch::kFloat64)
    );

    reduce_kernel<BLOCK_SIZE><<<num_blocks, BLOCK_SIZE>>>(
        data.data_ptr<float>(),
        partial.data_ptr<double>(),
        n
    );

    final_reduce_kernel<BLOCK_SIZE><<<1, BLOCK_SIZE>>>(
        partial.data_ptr<double>(),
        output.data_ptr<float>(),
        num_blocks
    );

    return output;
}
"""

cpp_source = "torch::Tensor vectorsum_cuda(torch::Tensor data, torch::Tensor output);"

_ext = load_inline(
    name="vectorsum",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["vectorsum_cuda"],
    extra_cuda_cflags=["-use_fast_math"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    data, output = data
    _ext.vectorsum_cuda(data, output)
    # reference returns a 0-d scalar from .sum(); collapse our {1} buffer to match
    return output.reshape(())