submission 772830

#!POPCORN leaderboard vectoradd_v2
#!POPCORN gpu A100

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


@triton.autotune(
    configs=[
        triton.Config({"BLOCK_SIZE": 256}, num_warps=4, num_stages=2),
        triton.Config({"BLOCK_SIZE": 512}, num_warps=4, num_stages=2),
        triton.Config({"BLOCK_SIZE": 1024}, num_warps=8, num_stages=3),
    ],
    key=["n_elements"],
)
@triton.jit
def add_kernel(a_ptr, b_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)

    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    # Load (coalesced)
    a = tl.load(a_ptr + offsets, mask=mask, other=0.0)
    b = tl.load(b_ptr + offsets, mask=mask, other=0.0)

    out = a + b

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


def custom_kernel(data: input_t) -> output_t:
    A, B, output = data
    n_elements = A.numel()

    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)

    add_kernel[grid](
        A,
        B,
        output,
        n_elements,
    )

    return output