submission 470769

#!POPCORN leaderboard vectoradd_v2
#!POPCORN gpu A100

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,
    N,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    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)


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