submission 750300

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


@triton.jit
def vecadd_kernel(
    A_ptr, B_ptr, C_ptr,
    N,
    BLOCK_SIZE: tl.constexpr,
):
    idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < N
    a = tl.load(A_ptr + idx, mask=mask)
    b = tl.load(B_ptr + idx, mask=mask)
    tl.store(C_ptr + idx, a + b, mask=mask)


def custom_kernel(data: input_t) -> output_t:
    A, B, C = data
    N = A.numel()
    BLOCK_SIZE = 1024
    grid = ((N + BLOCK_SIZE - 1) // BLOCK_SIZE,)
    vecadd_kernel[grid](A, B, C, N, BLOCK_SIZE=BLOCK_SIZE)
    return C