submission 66724

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


@triton.jit
def add_fp16_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    x = tl.load(x_ptr + offsets, mask=mask, other=0)
    y = tl.load(y_ptr + offsets, mask=mask, other=0)

    # For extra numerical headroom, you could do:
    # z = tl.cast(tl.cast(x, tl.float32) + tl.cast(y, tl.float32), tl.float16)
    z = x + y

    tl.store(out_ptr + offsets, z, mask=mask)


# User kernel implementation.
def custom_kernel(data: input_t) -> output_t:
    A, B, output = data
    n = A.numel()
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
    add_fp16_kernel[grid](
        A,
        B,
        output,
        n,
        BLOCK_SIZE=1024,  # try 1024, 2048, 4096
        num_warps=4,  # try 4, 8 depending on your GPU
    )
    return output