submission 515788

#!POPCORN leaderboard vectoradd_v2
#!POPCORN gpu A100

import torch
import triton
import triton.language as tl

@triton.jit
def vec_add_kernel(A_ptr, B_ptr, C_ptr, SIZE, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < SIZE
    a = tl.load(A_ptr + offs, mask=mask)
    b = tl.load(B_ptr + offs, mask=mask)
    tl.store(C_ptr + offs, a + b, mask=mask)

def custom_kernel(data):
    A, B, C = data
    SIZE = A.numel()
    BLOCK = 2048
    grid = (triton.cdiv(SIZE, BLOCK),)
    vec_add_kernel[grid](
        A.view(-1), B.view(-1), C.view(-1),
        SIZE, BLOCK
    )
    return C