submission 651140

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpu B200

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

CUDA_SRC = """// Your CUDA kernel and C++ launcher go here
#include <torch/extension.h>


__global__ void vec_sum(const float4* A, float* out, int N) {
    extern __shared__ float sdata[];
    int stride = gridDim.x * blockDim.x;
    float sum{};

    for (int idx = blockIdx.x * blockDim.x + threadIdx.x ; idx < N / 4 ; idx += stride) {
        float4 a = A[idx];
        sum += a.x + a.y + a.z + a.w;
    }

    sdata[threadIdx.x] = sum;
    __syncthreads();

    // Shared memory reduction until we're down to one warp
    for (int offset = blockDim.x / 2; offset >= 32; offset /= 2) {
        if (threadIdx.x < offset) {
            sdata[threadIdx.x] += sdata[threadIdx.x + offset];
        }
        __syncthreads();
    }

    // Warp-level reduction (no sync needed, threads execute in lockstep)
    if (threadIdx.x < 32) {
        float val = sdata[threadIdx.x];
        val += __shfl_down_sync(0xFFFFFFFF, val, 16);
        val += __shfl_down_sync(0xFFFFFFFF, val, 8);
        val += __shfl_down_sync(0xFFFFFFFF, val, 4);
        val += __shfl_down_sync(0xFFFFFFFF, val, 2);
        val += __shfl_down_sync(0xFFFFFFFF, val, 1);
        if (threadIdx.x == 0) {
            atomicAdd(out, val);
        }
    }
}

__global__ void vec_sum_tail(const float* A, float* out, int start, int N) {
    int idx = start + threadIdx.x;
    if (idx < N) {
        atomicAdd(out, A[idx]);
    }
}




torch::Tensor& vecsum(const torch::Tensor& in, torch::Tensor& out) {
    out.zero_();

    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, 0);

    int N = in.numel();
    int threads{256};
    int shared_mem_size = threads * sizeof(float);

    int blocks_per_SM;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks_per_SM, vec_sum, threads, shared_mem_size);

    int blocks = prop.multiProcessorCount * blocks_per_SM;

    int tail_start = (N / 4) * 4;
    int tail_count = N - tail_start;
    if (tail_count > 0) {
        vec_sum_tail<<<1, 256, 256 * sizeof(float)>>>(
            in.data_ptr<float>(),
            out.data_ptr<float>(),
            tail_start, N
        );
    }

    vec_sum<<<blocks, threads, threads * sizeof(float)>>>(
        reinterpret_cast<const float4*>(in.data_ptr<float>()),
        out.data_ptr<float>(),
        N
    );
    return out;
}
"""

CPP_SRC = """// Your C++ function declarations go here
torch::Tensor& vecsum(const torch::Tensor& in, torch::Tensor& out);
"""

module = load_inline(
    name='vecsum_module',
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=['vecsum'],
    verbose=True,
    extra_cuda_cflags=['-arch=sm_100', '--use_fast_math'],
)

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    return module.vecsum(data, output)[0]