submission 506787

#!POPCORN leaderboard vectorsum_v2

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

# EVOLVE-BLOCK-START
_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

// ---------------------------------------------------------------------------
// Warp-level reduction using shuffle (double precision)
// ---------------------------------------------------------------------------
__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;
}

// ---------------------------------------------------------------------------
// Phase 1: Each block computes a partial sum into partials[blockIdx.x]
//   - 4x unrolled float4 loads for ILP
//   - Warp shuffle + shared mem reduction
//   - No atomics – just a store per block
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(256, 4)
vector_sum_phase1(
    const float* __restrict__ x,
    double*      __restrict__ partials,
    const int n)
{
    extern __shared__ double sdata[];

    const int tid    = threadIdx.x;
    const int gid    = blockIdx.x * blockDim.x + tid;
    const int stride = gridDim.x * blockDim.x;

    double acc = 0.0;

    // --- Vectorized float4 loads with 4x unroll for ILP ---
    const int n_vec = n >> 2;  // n / 4
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x);

    // 4x unrolled loop
    const int stride4 = stride * 4;
    int i = gid;
    int limit = n_vec - 3 * stride;
    for (; i < limit; i += stride4) {
        float4 v0 = __ldg(&x4[i]);
        float4 v1 = __ldg(&x4[i + stride]);
        float4 v2 = __ldg(&x4[i + 2 * stride]);
        float4 v3 = __ldg(&x4[i + 3 * stride]);
        acc += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
        acc += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
        acc += (double)v2.x + (double)v2.y + (double)v2.z + (double)v2.w;
        acc += (double)v3.x + (double)v3.y + (double)v3.z + (double)v3.w;
    }
    // Cleanup remaining float4s
    for (; i < n_vec; i += stride) {
        float4 v = __ldg(&x4[i]);
        acc += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
    }

    // --- Tail elements ---
    const int tail_base = n_vec << 2;
    for (int j = tail_base + gid; j < n; j += stride)
        acc += (double)__ldg(&x[j]);

    // --- Warp-level reduce ---
    acc = warp_reduce_sum(acc);

    const int lane    = tid & 31;
    const int warp_id = tid >> 5;

    if (lane == 0)
        sdata[warp_id] = acc;
    __syncthreads();

    // --- Block-level reduce (first warp) ---
    const int n_warps = blockDim.x >> 5;  // 256/32 = 8
    acc = (tid < n_warps) ? sdata[tid] : 0.0;
    if (warp_id == 0)
        acc = warp_reduce_sum(acc);

    if (tid == 0)
        partials[blockIdx.x] = acc;
}

// ---------------------------------------------------------------------------
// Phase 2: Reduce partials array (small – just num_blocks doubles)
// Single block, single warp is enough for up to ~1024 partials
// ---------------------------------------------------------------------------
__global__ void vector_sum_phase2(
    const double* __restrict__ partials,
    float*        __restrict__ out,
    const int num_partials)
{
    double acc = 0.0;
    for (int i = threadIdx.x; i < num_partials; i += blockDim.x)
        acc += partials[i];

    acc = warp_reduce_sum(acc);

    // If more than 1 warp, use shared mem
    __shared__ double sdata[8];
    const int lane = threadIdx.x & 31;
    const int warp_id = threadIdx.x >> 5;

    if (lane == 0)
        sdata[warp_id] = acc;
    __syncthreads();

    const int n_warps = blockDim.x >> 5;
    if (warp_id == 0) {
        acc = (threadIdx.x < n_warps) ? sdata[threadIdx.x] : 0.0;
        acc = warp_reduce_sum(acc);
        if (threadIdx.x == 0)
            *out = (float)acc;
    }
}

// ---------------------------------------------------------------------------
// Host launcher
// ---------------------------------------------------------------------------
torch::Tensor vector_sum_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.is_contiguous());
    TORCH_CHECK(x.scalar_type() == torch::kFloat32);

    const int n = x.numel();

    // Persistent thread config: fill A100's 108 SMs with 4 blocks/SM
    const int THREADS = 256;
    const int BLOCKS  = 432;  // 108 SMs * 4
    const int n_warps = THREADS / 32;  // 8
    const size_t smem = n_warps * sizeof(double);

    // Allocate partials buffer and output
    auto partials = torch::empty({BLOCKS}, x.options().dtype(torch::kFloat64));
    auto out = torch::empty({1}, x.options().dtype(torch::kFloat32));

    vector_sum_phase1<<<BLOCKS, THREADS, smem>>>(
        x.data_ptr<float>(),
        partials.data_ptr<double>(),
        n);

    // Phase 2: reduce 432 doubles with 256 threads (1 block)
    vector_sum_phase2<<<1, 256, 8 * sizeof(double)>>>(
        partials.data_ptr<double>(),
        out.data_ptr<float>(),
        BLOCKS);

    return out[0];
}
"""

_cpp_src = r"""
torch::Tensor vector_sum_cuda(torch::Tensor x);
"""

_ext = load_inline(
    name="vectorsum_v2_opt",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["vector_sum_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo", "--ptxas-options=-v"],
    verbose=False,
)


def _custom_kernel(data: input_t) -> output_t:
    x, _ = data
    return _ext.vector_sum_cuda(x)


custom_kernel = torch.compile(_custom_kernel, mode="reduce-overhead")
# EVOLVE-BLOCK-END