submission 134798

#!POPCORN leaderboard vectorsum_v2

# This is a submission template for popcorn leaderboard 'vectorsum_v2'.
# Your task is as follows:
# > Implement a vector sum reduction kernel. This kernel computes the sum of all elements in the input tensor.
# > 
# > Input: A tensor of shape `(N,)` with values from a normal distribution with mean 0 and variance 1.
# > Output: A scalar value equal to the sum of all elements in the input tensor.
# The deadline for this leaderboard is 2025-12-30 00:00:00+00:00

# You can automatically route this file to specific GPUs by adding a line
# `#!POPCORN gpus <GPUs>` to the header of this file.
# Happy hacking!

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


cpp_source = """
#include <torch/extension.h>
#include <torch/torch.h>


torch::Tensor pad_to_multiple_of_4(torch::Tensor a) {
    // 获取当前 tensor 的大小
    int64_t current_size = a.size(0);
    if (current_size % 4 == 0) return a;

    // 计算到最近的 4 的倍数的差值
    int64_t pad_size = (4 - (current_size % 4)) % 4;

    // 创建一个填充的零张量，大小为需要填充的数量
    torch::Tensor padding = torch::zeros({pad_size}, a.options());

    // 拼接原始张量和填充的零张量
    return torch::cat({a, padding}, 0);
}


void vec_sum_kernel_launcher( float* a, double* y, int N);

torch::Tensor vector_sum(torch::Tensor a) {
    a = pad_to_multiple_of_4(a);
    int64_t N = a.numel();
    // Use double precision for accumulation, then convert to float32 to match reference
    torch::Tensor y = torch::tensor(0.0, torch::dtype(torch::kFloat64)).to(a.device());

    vec_sum_kernel_launcher(
        a.data_ptr<float>(),
        y.data_ptr<double>(),
        (int)N
    );

    // Convert to float32 to match reference output type
    return y.to(torch::kFloat32);
}
"""

cuda_source = """
#include <cuda_runtime.h>

#define WARP_SIZE 32
#define FLOAT4(value) (reinterpret_cast<float4 *>(&(value))[0])

template < int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_sum_f32(float val) {
#pragma unroll
  for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {
    val += __shfl_xor_sync(0xffffffff, val, mask);
  }
  return val;
}

template <const int NUM_THREADS = 256/4>
__global__ void vec_sum_kernel(float* a, double* y, int N) {
  int tid = threadIdx.x;
  int idx = (blockIdx.x * NUM_THREADS + tid) * 4;
  constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
  __shared__ float reduce_smem[NUM_WARPS];

  float4 reg_a = FLOAT4(a[idx]);
  float sum = (idx < N) ? (reg_a.x + reg_a.y + reg_a.z + reg_a.w) : 0.0f;
  int warp = tid / WARP_SIZE;
  int lane = tid % WARP_SIZE;
  // perform warp sync reduce.
  sum = warp_reduce_sum_f32<WARP_SIZE>(sum);
  // warp leaders store the data to shared memory.
  if (lane == 0)
    reduce_smem[warp] = sum;
  __syncthreads(); // make sure the data is in shared memory.
  // the first warp compute the final sum.
  sum = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;
  if (warp == 0)
    sum = warp_reduce_sum_f32<NUM_WARPS>(sum);
  if (tid == 0)
    atomicAdd(y, sum);
}

void vec_sum_kernel_launcher(float* a, double* y, int N) {
    dim3 block = 512;  
    dim3 grid((N + block.x * 4 - 1) / (block.x * 4));  // Dynamic grid size
    vec_sum_kernel<512><<<grid, block>>>(a, y, N);
}
"""

module = load_inline(
    name="vecsum",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["vector_sum"],  # this exposes the function to Python
    extra_cuda_cflags=[
        "-std=c++17",
        # "-gencode=arch=compute_100a,code=sm_100a",
        # "--ptxas-options=--gpu-name=sm_100a",
        "-O3",
        "-w",
        "-maxrregcount=32",
        "--use_fast_math",
        "-allow-unsupported-compiler",
    ],
    extra_ldflags=["-lcuda", "-lcublas"],
    verbose=True,
)

def custom_kernel(data: input_t) -> output_t:
    return module.vector_sum(data[0])
    

# a = torch.randn(1024, device="cuda")
# print(a.sum())
# y = module.vector_sum(a)
# print(a.sum())
# print(y)