submission 681553

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpu B200

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

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

// Persistent state for partials and counter
static double* d_partials = nullptr;
static unsigned int* d_counter = nullptr;
static int partials_capacity = 0;

// Single-kernel reduction using last-block-does-final technique.
// Each block reduces its chunk, writes partial to global mem, atomicInc a counter.
// The last block to finish does the final reduction of all partials.
__global__ void vectorsum_lastblock(const float* __restrict__ data, float* __restrict__ output,
                                     double* __restrict__ partials, unsigned int* __restrict__ counter,
                                     int n, int num_blocks_total) {
    extern __shared__ double sdata[];

    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int global_tid = bid * blockDim.x + tid;
    const int grid_stride = gridDim.x * blockDim.x;
    const int warp_id = tid / 32;
    const int lane = tid % 32;
    const int num_warps = blockDim.x / 32;

    // Phase 1: Each thread accumulates its chunk in float64
    double sum = 0.0;

    const float4* data4 = reinterpret_cast<const float4*>(data);
    const int n4 = n / 4;

    // 8x unrolled float4 loads
    int idx = global_tid;
    while (idx + 7 * grid_stride < n4) {
        float4 v0 = __ldg(&data4[idx]);
        float4 v1 = __ldg(&data4[idx + grid_stride]);
        float4 v2 = __ldg(&data4[idx + 2 * grid_stride]);
        float4 v3 = __ldg(&data4[idx + 3 * grid_stride]);
        float4 v4 = __ldg(&data4[idx + 4 * grid_stride]);
        float4 v5 = __ldg(&data4[idx + 5 * grid_stride]);
        float4 v6 = __ldg(&data4[idx + 6 * grid_stride]);
        float4 v7 = __ldg(&data4[idx + 7 * grid_stride]);

        sum += (double)v0.x + (double)v0.y + (double)v0.z + (double)v0.w;
        sum += (double)v1.x + (double)v1.y + (double)v1.z + (double)v1.w;
        sum += (double)v2.x + (double)v2.y + (double)v2.z + (double)v2.w;
        sum += (double)v3.x + (double)v3.y + (double)v3.z + (double)v3.w;
        sum += (double)v4.x + (double)v4.y + (double)v4.z + (double)v4.w;
        sum += (double)v5.x + (double)v5.y + (double)v5.z + (double)v5.w;
        sum += (double)v6.x + (double)v6.y + (double)v6.z + (double)v6.w;
        sum += (double)v7.x + (double)v7.y + (double)v7.z + (double)v7.w;

        idx += 8 * grid_stride;
    }
    while (idx < n4) {
        float4 v = __ldg(&data4[idx]);
        sum += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
        idx += grid_stride;
    }
    for (int i = n4 * 4 + global_tid; i < n; i += grid_stride) {
        sum += (double)__ldg(&data[i]);
    }

    // Warp-level reduction
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
    }

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

    if (warp_id == 0) {
        sum = (lane < num_warps) ? sdata[lane] : 0.0;
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
        }
    }

    // Block leader writes partial and increments counter
    __shared__ bool is_last_block;
    if (tid == 0) {
        partials[bid] = sum;
        __threadfence();
        unsigned int old = atomicInc(counter, num_blocks_total - 1);
        is_last_block = (old == (unsigned int)(num_blocks_total - 1));
    }
    __syncthreads();

    // Phase 2: The last block reduces all partials
    if (is_last_block) {
        double phase2_sum = 0.0;
        for (int i = tid; i < num_blocks_total; i += blockDim.x) {
            phase2_sum += partials[i];
        }

        for (int offset = 16; offset > 0; offset >>= 1) {
            phase2_sum += __shfl_down_sync(0xFFFFFFFF, phase2_sum, offset);
        }

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

        if (warp_id == 0) {
            phase2_sum = (lane < num_warps) ? sdata[lane] : 0.0;
            for (int offset = 16; offset > 0; offset >>= 1) {
                phase2_sum += __shfl_down_sync(0xFFFFFFFF, phase2_sum, offset);
            }
            if (lane == 0) {
                *output = (float)phase2_sum;
                *counter = 0;
            }
        }
    }
}

void vectorsum_raw(int64_t data_ptr, int64_t output_ptr, int N) {
    const int block_size = 256;
    const int shared_mem = (block_size / 32) * sizeof(double);

    int num_blocks = 148 * 4;  // 592 blocks for 148 SMs
    int blocks_needed = (N / 4 + block_size - 1) / block_size;
    if (num_blocks > blocks_needed) num_blocks = blocks_needed;
    if (num_blocks < 1) num_blocks = 1;

    // Allocate partials and counter (cached)
    if (d_partials == nullptr || partials_capacity < num_blocks) {
        if (d_partials) cudaFree(d_partials);
        if (d_counter) cudaFree(d_counter);
        cudaMalloc(&d_partials, num_blocks * sizeof(double));
        cudaMalloc(&d_counter, sizeof(unsigned int));
        cudaMemset(d_counter, 0, sizeof(unsigned int));
        partials_capacity = num_blocks;
    }

    vectorsum_lastblock<<<num_blocks, block_size, shared_mem>>>(
        reinterpret_cast<const float*>(data_ptr),
        reinterpret_cast<float*>(output_ptr),
        d_partials,
        d_counter,
        N,
        num_blocks
    );
}
"""

cpp_src = r"""
void vectorsum_raw(int64_t data_ptr, int64_t output_ptr, int N);
"""

_ext = load_inline(
    name="vectorsum_lastblock",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["vectorsum_raw"],
    with_cuda=True,
    extra_cflags=["-O3", "-std=c++17"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17",
                       "-gencode=arch=compute_100,code=sm_100"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    data_tensor, output_tensor = data
    _ext.vectorsum_raw(data_tensor.data_ptr(), output_tensor.data_ptr(), data_tensor.numel())
    return output_tensor[0]