submission 780483

import torch
from torch.utils.cpp_extension import load_inline

# ─────────────────────────────────────────────
# CUDA kernel
# ─────────────────────────────────────────────
cuda_src = r'''
#include <cuda_runtime.h>

// Shuffle a double across warp lanes by splitting into two 32-bit halves.
// Required because __shfl_down_sync only natively handles 32-bit values
// on older architectures (sm_60 handles double natively, but this is portable).
__device__ __forceinline__ double shfl_down_double(double val, int offset) {
    unsigned int lo = __double2loint(val);
    unsigned int hi = __double2hiint(val);
    lo = __shfl_down_sync(0xffffffff, lo, offset);
    hi = __shfl_down_sync(0xffffffff, hi, offset);
    return __hiloint2double(hi, lo);
}

__global__ void sum_kernel(const float* __restrict__ input, double* output, int n) {
    double acc = 0.0;

    // Grid-stride loop: each thread accumulates multiple elements in double.
    // This handles any N regardless of grid size, and keeps the grid small
    // so we don't launch 200k blocks for large inputs.
    int idx    = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x  * blockDim.x;
    for (int i = idx; i < n; i += stride)
        acc += (double)input[i];

    // ── Warp-level reduction (no shared memory needed) ──────────────────────
    for (int offset = 16; offset > 0; offset >>= 1)
        acc += shfl_down_double(acc, offset);

    // ── Block-level reduction via shared memory (one slot per warp) ─────────
    __shared__ double smem[32];          // max 32 warps per block (1024 threads)
    int lane   = threadIdx.x & 31;
    int warpid = threadIdx.x >> 5;
    int nwarps = (blockDim.x + 31) >> 5;

    if (lane == 0) smem[warpid] = acc;
    __syncthreads();

    // First warp finishes the block reduction and atomically adds to global output
    if (warpid == 0) {
        acc = (lane < nwarps) ? smem[lane] : 0.0;
        for (int offset = 16; offset > 0; offset >>= 1)
            acc += shfl_down_double(acc, offset);
        if (lane == 0)
            atomicAdd(output, acc);      // double atomicAdd: native on sm_60+
    }
}

// Host-side launcher called from the C++ wrapper below
void launch_sum_kernel(const float* input, double* output, int n) {
    const int threads = 256;
    // Cap blocks so we don't over-launch; grid-stride loop handles the rest.
    // 2048 blocks * 256 threads = 524288 threads. For 52M elements each
    // thread handles ~100 elements — good occupancy/memory-bandwidth balance.
    const int blocks = min((n + threads - 1) / threads, 2048);
    sum_kernel<<<blocks, threads>>>(input, output, n);
}
'''

# ─────────────────────────────────────────────
# PyTorch / C++ wrapper
# ─────────────────────────────────────────────
cpp_src = r'''
#include <torch/extension.h>

void launch_sum_kernel(const float* input, double* output, int n);

torch::Tensor fast_sum(torch::Tensor input) {
    TORCH_CHECK(input.is_cuda(),  "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == torch::kFloat32, "input must be float32");
    TORCH_CHECK(input.is_contiguous(), "input must be contiguous");

    int n = input.numel();

    // Accumulator lives in double — same as the reference implementation
    auto acc = torch::zeros(
        {1},
        torch::TensorOptions()
            .dtype(torch::kDouble)
            .device(input.device())
    );

    launch_sum_kernel(
        input.data_ptr<float>(),
        acc.data_ptr<double>(),
        n
    );

    // Cast back to float32 to match reference output type
    return acc.to(torch::kFloat32);
}
'''

# ─────────────────────────────────────────────
# Compile
# ─────────────────────────────────────────────
_module = load_inline(
    name='fast_sum_v1',
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=['fast_sum'],
    with_cuda=True,
    extra_cuda_cflags=['-O3', '--use_fast_math'],
)

# ─────────────────────────────────────────────
# Entry point called by the autograder
# ─────────────────────────────────────────────
def custom_kernel(data):
    input_tensor, _ = data              # unpack (input, empty output buffer)
    result = _module.fast_sum(input_tensor)
    return result.squeeze()             # return 0-dim float32 scalar tensor