submission 665622

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

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

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

__global__ void fast_sum_kernel(
    const float* __restrict__ x,
    double* __restrict__ partial,
    int n
) {
    double acc = 0.0;
    int tid    = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    const float4* x4 = reinterpret_cast<const float4*>(x);
    int n4 = n >> 2;

    // 4x unrolled float4 grid-stride loop
    int i = tid;
    for (; i + stride * 3 < n4; i += stride * 4) {
        float4 a = x4[i];
        float4 b = x4[i + stride];
        float4 c = x4[i + stride * 2];
        float4 d = x4[i + stride * 3];
        acc += (double)a.x + (double)a.y + (double)a.z + (double)a.w;
        acc += (double)b.x + (double)b.y + (double)b.z + (double)b.w;
        acc += (double)c.x + (double)c.y + (double)c.z + (double)c.w;
        acc += (double)d.x + (double)d.y + (double)d.z + (double)d.w;
    }
    for (; i < n4; i += stride) {
        float4 v = x4[i];
        acc += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
    }
    // Scalar tail
    for (int j = n4 * 4 + tid; j < n; j += blockDim.x * gridDim.x)
        acc += (double)x[j];

    // Warp reduce
    acc = warp_reduce_sum(acc);

    __shared__ double shared[32];
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    if (lane == 0) shared[wid] = acc;
    __syncthreads();

    if (wid == 0) {
        int nwarps = blockDim.x >> 5;
        acc = (lane < nwarps) ? shared[lane] : 0.0;
        acc = warp_reduce_sum(acc);
        if (lane == 0) partial[blockIdx.x] = acc;
    }
}

torch::Tensor fast_sum(torch::Tensor x, torch::Tensor partial) {
    int n = x.numel();
    const int threads = 1024;
    const int blocks  = 512;

    fast_sum_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        partial.data_ptr<double>(),
        n
    );

    // Sum partials and return fresh tensor — never reuse output across calls
    return partial.sum().to(torch::kFloat32);
}
"""

_cpp_src = "torch::Tensor fast_sum(torch::Tensor x, torch::Tensor partial);"

_ext = load_inline(
    name="fast_sum_ext3",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fast_sum"],
    with_cuda=True,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

_partial = torch.empty(512, device="cuda", dtype=torch.float64)


def custom_kernel(data: input_t) -> output_t:
    x, _ = data
    return _ext.fast_sum(x, _partial)