submission 681746

#!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>

static float* d_partials = nullptr;
static unsigned int* d_counter = nullptr;
static int partials_capacity = 0;

// Ultra-low occupancy: 2 blocks/SM = 296 blocks on 148 SMs
// Maximum registers per thread, each thread processes massive chunk
__global__ __launch_bounds__(256, 2)
void vectorsum_fast(const float* __restrict__ data, float* __restrict__ output,
                    float* __restrict__ partials, unsigned int* __restrict__ counter) {
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int global_tid = bid * 256 + tid;
    const int grid_stride = 296 * 256;
    const int lane = tid & 31;
    const int warp_id = tid >> 5;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float s4 = 0.f, s5 = 0.f, s6 = 0.f, s7 = 0.f;

    const float4* data4 = reinterpret_cast<const float4*>(data);

    int idx = global_tid;
    #pragma unroll 4
    while (idx + grid_stride < 13107200) {
        const float4* p0 = &data4[idx];
        const float4* p1 = &data4[idx + grid_stride];
        float4 v0, v1;
        asm volatile("ld.global.cs.v4.f32 {%0,%1,%2,%3}, [%4];" : "=f"(v0.x), "=f"(v0.y), "=f"(v0.z), "=f"(v0.w) : "l"(p0));
        asm volatile("ld.global.cs.v4.f32 {%0,%1,%2,%3}, [%4];" : "=f"(v1.x), "=f"(v1.y), "=f"(v1.z), "=f"(v1.w) : "l"(p1));
        s0 += v0.x; s1 += v0.y; s2 += v0.z; s3 += v0.w;
        s4 += v1.x; s5 += v1.y; s6 += v1.z; s7 += v1.w;
        idx += grid_stride * 2;
    }
    if (idx < 13107200) {
        const float4* p0 = &data4[idx];
        float4 v0;
        asm volatile("ld.global.cs.v4.f32 {%0,%1,%2,%3}, [%4];" : "=f"(v0.x), "=f"(v0.y), "=f"(v0.z), "=f"(v0.w) : "l"(p0));
        s0 += v0.x; s1 += v0.y; s2 += v0.z; s3 += v0.w;
    }

    float sum = (s0 + s1) + (s2 + s3) + (s4 + s5) + (s6 + s7);

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);

    __shared__ float smem[8];
    if (lane == 0) smem[warp_id] = sum;
    __syncthreads();

    if (warp_id == 0) {
        sum = (lane < 8) ? smem[lane] : 0.f;
        #pragma unroll
        for (int offset = 4; offset > 0; offset >>= 1)
            sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
    }

    __shared__ bool is_last;
    if (tid == 0) {
        partials[bid] = sum;
        __threadfence();
        unsigned int old = atomicInc(counter, 295);
        is_last = (old == 295);
    }
    __syncthreads();

    if (is_last) {
        float psum = 0.f;
        #pragma unroll 2
        for (int i = tid; i < 296; i += 256)
            psum += partials[i];
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1)
            psum += __shfl_down_sync(0xFFFFFFFF, psum, offset);
        if (lane == 0) smem[warp_id] = psum;
        __syncthreads();
        if (warp_id == 0) {
            psum = (lane < 8) ? smem[lane] : 0.f;
            #pragma unroll
            for (int offset = 4; offset > 0; offset >>= 1)
                psum += __shfl_down_sync(0xFFFFFFFF, psum, offset);
            if (lane == 0) *output = psum;
        }
    }
}

// Generic for test sizes
__global__ __launch_bounds__(256, 8)
void vectorsum_generic(const float* __restrict__ data, float* __restrict__ output,
                       float* __restrict__ partials, unsigned int* __restrict__ counter,
                       int n) {
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int global_tid = bid * 256 + tid;
    const int grid_stride = gridDim.x * 256;
    const int lane = tid & 31;
    const int warp_id = tid >> 5;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float s4 = 0.f, s5 = 0.f, s6 = 0.f, s7 = 0.f;

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

    int idx = global_tid;
    while (idx + grid_stride < n4) {
        const float4* p0 = &data4[idx];
        const float4* p1 = &data4[idx + grid_stride];
        float4 v0, v1;
        asm volatile("ld.global.cs.v4.f32 {%0,%1,%2,%3}, [%4];" : "=f"(v0.x), "=f"(v0.y), "=f"(v0.z), "=f"(v0.w) : "l"(p0));
        asm volatile("ld.global.cs.v4.f32 {%0,%1,%2,%3}, [%4];" : "=f"(v1.x), "=f"(v1.y), "=f"(v1.z), "=f"(v1.w) : "l"(p1));
        s0 += v0.x; s1 += v0.y; s2 += v0.z; s3 += v0.w;
        s4 += v1.x; s5 += v1.y; s6 += v1.z; s7 += v1.w;
        idx += grid_stride * 2;
    }
    if (idx < n4) {
        const float4* p0 = &data4[idx];
        float4 v0;
        asm volatile("ld.global.cs.v4.f32 {%0,%1,%2,%3}, [%4];" : "=f"(v0.x), "=f"(v0.y), "=f"(v0.z), "=f"(v0.w) : "l"(p0));
        s0 += v0.x; s1 += v0.y; s2 += v0.z; s3 += v0.w;
        idx += grid_stride;
    }
    for (int i = n4 * 4 + global_tid; i < n; i += grid_stride) {
        float val;
        asm volatile("ld.global.cs.f32 %0, [%1];" : "=f"(val) : "l"(&data[i]));
        s0 += val;
    }

    float sum = (s0 + s1) + (s2 + s3) + (s4 + s5) + (s6 + s7);

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);

    __shared__ float smem[8];
    if (lane == 0) smem[warp_id] = sum;
    __syncthreads();

    if (warp_id == 0) {
        sum = (lane < 8) ? smem[lane] : 0.f;
        #pragma unroll
        for (int offset = 4; offset > 0; offset >>= 1)
            sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
    }

    __shared__ bool is_last;
    if (tid == 0) {
        partials[bid] = sum;
        __threadfence();
        unsigned int old = atomicInc(counter, gridDim.x - 1);
        is_last = (old == gridDim.x - 1);
    }
    __syncthreads();

    if (is_last) {
        float psum = 0.f;
        for (int i = tid; i < (int)gridDim.x; i += 256)
            psum += partials[i];
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1)
            psum += __shfl_down_sync(0xFFFFFFFF, psum, offset);
        if (lane == 0) smem[warp_id] = psum;
        __syncthreads();
        if (warp_id == 0) {
            psum = (lane < 8) ? smem[lane] : 0.f;
            #pragma unroll
            for (int offset = 4; offset > 0; offset >>= 1)
                psum += __shfl_down_sync(0xFFFFFFFF, psum, offset);
            if (lane == 0) *output = psum;
        }
    }
}

void vectorsum_raw(int64_t data_ptr, int64_t output_ptr, int N) {
    int num_blocks;

    if (N == 52428800) {
        num_blocks = 296;
    } else {
        num_blocks = 1184;
    }

    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(float));
        cudaMalloc(&d_counter, sizeof(unsigned int));
        cudaMemset(d_counter, 0, sizeof(unsigned int));
        partials_capacity = num_blocks;
    }

    if (N == 52428800) {
        vectorsum_fast<<<296, 256, 0>>>(
            reinterpret_cast<const float*>(data_ptr),
            reinterpret_cast<float*>(output_ptr),
            d_partials, d_counter);
    } else {
        int blocks_needed = (N / 4 + 255) / 256;
        int nb = (blocks_needed < num_blocks) ? blocks_needed : num_blocks;
        if (nb < 1) nb = 1;
        vectorsum_generic<<<nb, 256, 0>>>(
            reinterpret_cast<const float*>(data_ptr),
            reinterpret_cast<float*>(output_ptr),
            d_partials, d_counter, N);
    }
}
"""

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

_ext = load_inline(
    name="vectorsum_296blk_popcorn",
    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]