submission 538949

from utils import make_match_reference, DeterministicContext
import torch
from task import input_t, output_t
import triton
import triton.language as tl

@triton.jit(
    version=None,
    repr=None,
    launch_metadata=None,
    do_not_specialize=None,
    do_not_specialize_on_alignment=None,
    debug=None,
    noinline=None,
)
def _vadd(
    p_v1: torch.Tensor,
    p_v2: torch.Tensor,
    p_o: torch.Tensor,
    n: int,
    tile_size_x: tl.constexpr,
):
    pidx = tl.program_id(axis=0)

    stride = pidx * tile_size_x + tl.arange(start=0, end=tile_size_x)
    mask = stride < n

    v_v1 = tl.load(pointer=p_v1 + stride, mask=mask, other=0.0)
    v_v2 = tl.load(pointer=p_v2 + stride, mask=mask, other=0.0)

    tl.store(pointer=p_o + stride, value=(v_v1 + v_v2), mask=mask)

def vadd(
    v1: torch.Tensor,
    v2: torch.Tensor,
    o: torch.Tensor
):
    n = v1.shape[0] * v1.shape[0]
    tile_size = 32768
    grid = (triton.cdiv(n, tile_size),)

    _vadd[grid](p_v1=v1, p_v2=v2, p_o=o, n=n, tile_size_x=tile_size)

    return o


def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of vector addition using PyTorch.
    Args:
        data: Tuple of tensors [A, B] to be added.
    Returns:
        Tensor containing element-wise sums.
    """
    with DeterministicContext():
        A, B, output = data
        output[...] = vadd(A, B, output)
        return output