submission 126194

# #!POPCORN leaderboard vectoradd_v2
# #!POPCORN gpus L4

# #!/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>
# #include <cuda_runtime.h>

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

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

# // 1D vector add: c[i] = a[i] + b[i] for arbitrary floating scalar_t
# template <typename scalar_t>
# __global__ void vecadd_kernel(const scalar_t* __restrict__ a,
#                               const scalar_t* __restrict__ b,
#                               scalar_t* __restrict__ c,
#                               long n) {
#     long i = blockIdx.x * blockDim.x + threadIdx.x;
#     if (i < n) {
#         c[i] = a[i] + b[i];
#     }
# }

# // Definition of vecadd_cuda declared in CPP_SRC
# torch::Tensor vecadd_cuda(torch::Tensor a, torch::Tensor b) {
#     // Assume inputs are already contiguous in this benchmark harness
#     auto a_c = a;
#     auto b_c = b;

#     TORCH_CHECK(a_c.sizes() == b_c.sizes(), "Input tensors must have same shape");

#     auto out = torch::empty_like(a_c);
#     long n = a_c.numel();

#     if (n == 0) {
#         return out;
#     }

#     constexpr int threads = 256;
#     // Enough threads to cover all elements, 1 element per thread
#     int blocks = static_cast<int>((n + threads - 1) / threads);

#     // Dispatch on the actual scalar type: float32, float16, bfloat16, double, ...
#     AT_DISPATCH_FLOATING_TYPES_AND_HALF(a_c.scalar_type(), "vecadd_cuda", [&] {
#         const scalar_t* a_ptr = a_c.data_ptr<scalar_t>();
#         const scalar_t* b_ptr = b_c.data_ptr<scalar_t>();
#         scalar_t* out_ptr     = out.data_ptr<scalar_t>();

#         vecadd_kernel<scalar_t><<<blocks, threads>>>(
#             a_ptr, b_ptr, out_ptr, n
#         );
#     });

#     return out;
# }
# """

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


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

#     The harness passes a tuple with at least two elements, e.g.:
#         (a, b, ...)

#     We:
#       - grab a and b,
#       - move them to CUDA (keeping their original dtype),
#       - run the CUDA vecadd,
#       - return a single tensor.
#     No checks, no dtype conversions.
#     """
#     # a = data[0]  # keep this to prevent error
#     # b = data[1]  # keep this

#     # Move to CUDA; keeps dtype, so we stay in orig_dtype
#     # a = a.cuda(non_blocking=True)
#     # b = b.cuda(non_blocking=True)

#     out = vecadd_mod.vecadd_cuda(data[0], data[1])

#     # No explicit synchronize: let the harness force sync when it needs the result
#     return out

#!POPCORN leaderboard vectoradd_v2
#!POPCORN gpus L4

#!/usr/bin/env python3
import torch
import triton
import triton.language as tl


@triton.jit
def vecadd_kernel(a_ptr, b_ptr, c_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    """
    Each Triton program (block) handles BLOCK_SIZE elements.
    """
    pid = tl.program_id(0)  # 1D grid
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    a = tl.load(a_ptr + offsets, mask=mask)
    b = tl.load(b_ptr + offsets, mask=mask)
    c = a + b
    tl.store(c_ptr + offsets, c, mask=mask)


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

    The harness passes a tuple with at least two elements, e.g.:
        (a, b, ...)

    We:
      - grab a and b,
      - move them to CUDA (keeping their original dtype),
      - run the Triton vecadd kernel,
      - return a single tensor.
    """
    a = data[0]  # keep this to prevent error
    b = data[1]  # keep this

    # Move to CUDA; keep original dtype
    a = a.cuda(non_blocking=True)
    b = b.cuda(non_blocking=True)

    # Shapes must match
    assert a.shape == b.shape, "Input tensors must have the same shape"

    # Ensure contiguous (cheap if already contiguous)
    a_c = a if a.is_contiguous() else a.contiguous()
    b_c = b if b.is_contiguous() else b.contiguous()

    out = torch.empty_like(a_c)
    n_elements = out.numel()

    if n_elements == 0:
        return out

    BLOCK_SIZE = 1024  # elements per Triton program
    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)

    # Launch Triton kernel
    vecadd_kernel[grid](
        a_c,
        b_c,
        out,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=4,
        num_stages=1,
    )

    return out