submission 34544

#!POPCORN leaderboard vectoradd_v2
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline


cuda_src = r"""
template <typename T>
__global__ void vectoradd_v2_kernel(const T* A, const T* B, T* C, long long n) {
    long long idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) C[idx] = A[idx] + B[idx];
}


void vectoradd_v2(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
    auto n = A.numel();
    int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);

    if (A.scalar_type() == at::kFloat) {
        vectoradd_v2_kernel<float><<<blocks, threads>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), n);
    } else if (A.scalar_type() == at::kHalf) {
        vectoradd_v2_kernel<at::Half><<<blocks, threads>>>(
            A.data_ptr<at::Half>(), B.data_ptr<at::Half>(), C.data_ptr<at::Half>(), n);
    } else if (A.scalar_type() == at::kDouble) {
        vectoradd_v2_kernel<double><<<blocks, threads>>>(
            A.data_ptr<double>(), B.data_ptr<double>(), C.data_ptr<double>(), n);
    } else {
        TORCH_CHECK(false, "Unsupported dtype");
    }
}
"""

cpp_src = "void vectoradd_v2(torch::Tensor A, torch::Tensor B, torch::Tensor C);"

module = load_inline(
    name="vector_add_module",
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=["vectoradd_v2"],
    with_cuda=True,
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

def inline_kernel(data: input_t) -> output_t:
    A, B, C = data
    module.vectoradd_v2(A, B, C)
    return C

def baseline_kernel(data: input_t) -> output_t:
    A, B, C = data
    C.copy_(A + B)
    return C

custom_kernel = inline_kernel

# import torch
# # Sanity check
# N = 1_000_000
# a = torch.randn(N, device="cuda", dtype=torch.float32)
# b = torch.randn_like(a)
# out = torch.empty_like(a)

# module.vectoradd_v2(a, b, out)   # launches the kernel on the current CUDA stream

# # correctness check
# torch.testing.assert_close(out, a + b)
# print("OK:", out[:3], (a[:3] + b[:3]))