submission 35770

#!POPCORN leaderboard vectoradd_v2

# This is a submission template for popcorn leaderboard 'vectoradd_v2'.
# Your task is as follows:
# > Implement a float16 vector addition kernel.
# > 
# > Input: tuple(torch.Tensor, torch.Tensor) with tensors of shape (N, N) and type torch.float16. These tensors are from
# > a normal distribution with mean 0 and variance 1.
# > Output: torch.Tensor of shape (N, N) and type torch.float16
# 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!

import triton
import triton.language as tl
from task import input_t, output_t

@triton.jit
def add_kernel(x_ptr, y_ptr, z_ptr, n, b_sz: tl.constexpr):
  pid = tl.program_id(axis=0)
  start = pid * b_sz
  offsets = start + tl.arange(0, b_sz)
  mask = offsets < n
  x = tl.load(x_ptr + offsets, mask=mask)
  y = tl.load(y_ptr + offsets, mask=mask)
  z = x + y
  tl.store(z_ptr + offsets, z, mask=mask)


def custom_kernel(data: input_t) -> output_t:
    x, y, z = data
    n = x.numel()
    grid_lam = lambda meta: (triton.cdiv(n, meta["b_sz"]), )

    add_kernel[grid_lam](x, y, z, n, b_sz=1024)

    return z