submission 126543

#!POPCORN leaderboard vectorsum_v2
#!POPCORN gpus H100

#!/usr/bin/env python3
import torch
from torch.utils.cpp_extension import load_inline

# 1) C++ side: declaration so Python can call into the CUDA implementation
CPP_SRC = r"""
#include <torch/extension.h>

// Forward declaration; definition will be in the CUDA translation unit.
torch::Tensor vectorsum_cuda(torch::Tensor a);
"""

# 2) CUDA side: reduction kernel + dtype dispatch
CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/ATen.h>

// 1D reduction: sum(a[i]) into a scalar, with per-block partials
template <typename scalar_t, typename acc_t>
__global__ void vectorsum_kernel(const scalar_t* __restrict__ a,
                                 acc_t* __restrict__ partial,
                                 long n) {
    extern __shared__ acc_t smem[];

    const int tid     = threadIdx.x;
    const int block_n = blockDim.x * gridDim.x;
    long idx          = blockIdx.x * blockDim.x + tid;

    // 1) Per-thread strided accumulation in higher-precision acc_t
    acc_t local = static_cast<acc_t>(0);
    for (; idx < n; idx += block_n) {
        local += static_cast<acc_t>(a[idx]);
    }

    // 2) Write to shared memory
    smem[tid] = local;
    __syncthreads();

    // 3) Block-wide reduction in shared memory
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            smem[tid] += smem[tid + s];
        }
        __syncthreads();
    }

    // 4) Thread 0 writes the block's partial sum
    if (tid == 0) {
        partial[blockIdx.x] = smem[0];
    }
}

// Definition of vectorsum_cuda declared in CPP_SRC
torch::Tensor vectorsum_cuda(torch::Tensor a) {
    auto a_c = a.contiguous();
    const long n = a_c.numel();

    if (n == 0) {
        // Handle empty input: return 0.0f on same device
        auto empty_out = torch::zeros({}, a_c.options().dtype(torch::kFloat32));
        return empty_out;
    }

    constexpr int threads = 256;
    int blocks = static_cast<int>((n + threads - 1) / threads);
    if (blocks < 1)   blocks = 1;
    if (blocks > 1024) blocks = 1024;  // Safety cap

    // Partial sums kept in float64 for better numerical agreement with reference
    auto partial = torch::empty({blocks}, a_c.options().dtype(torch::kFloat64));

    AT_DISPATCH_FLOATING_TYPES_AND_HALF(a_c.scalar_type(), "vectorsum_cuda", [&] {
        using acc_t = double;
        vectorsum_kernel<scalar_t, acc_t>
            <<<blocks, threads, threads * sizeof(acc_t)>>>(
                a_c.data_ptr<scalar_t>(),
                partial.data_ptr<acc_t>(),
                n);
    });

    // Final reduction on GPU in float64, then cast to float32
    auto total64 = partial.sum();                       // float64 scalar, same device as a_c
    auto total32 = total64.to(torch::kFloat32);         // match reference dtype

    return total32;
}
"""

# 3) Build extension once at import time.
vectorsum_mod = load_inline(
    name="vectorsum_ext",
    cpp_sources=CPP_SRC,
    cuda_sources=CUDA_SRC,
    functions=["vectorsum_cuda"],
    with_cuda=True,
    verbose=False,
)


@torch.inference_mode()
def custom_kernel(data):
    """
    GPU Mode entrypoint.

    The harness passes (input_tensor, output_tensor).

    We:
      - grab the input tensor,
      - ensure it's on CUDA,
      - call the CUDA vectorsum kernel,
      - return a single scalar tensor (float32 on CUDA).
    """
    x, out_buf = data  # out_buf is preallocated but we don't need to use it

    # Ensure tensor is on CUDA; generator already does this, but be robust
    x = x.cuda(non_blocking=True)

    out = vectorsum_mod.vectorsum_cuda(x)

    # Could also do out_buf.copy_(out) and return out_buf, but returning `out` is fine
    return out