submission 66933

#!POPCORN leaderboard vectoradd_v2

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


@triton.jit
def add_kernel(A_ptr, B_ptr, C_ptr, size, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < size
    A = tl.load(A_ptr + offsets, mask=mask)
    B = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, A + B, mask=mask)


def custom_kernel(data: input_t) -> output_t:
    A, B, C = data
    A, B, C = A.contiguous(), B.contiguous(), C.contiguous()
    size = A.numel()
    BLOCK_SIZE = 1024
    grid = (triton.cdiv(size, BLOCK_SIZE),)
    add_kernel[grid](A, B, C, size, BLOCK_SIZE=BLOCK_SIZE)
    return C