submission 67559

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


# Fixed optimal config - no autotuning variance
@triton.jit
def vecadd_fp16_kernel(A, B, C, N, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < N

    # Load
    a = tl.load(A + offs, mask=mask, other=0.0)
    b = tl.load(B + offs, mask=mask, other=0.0)

    # Compute
    c = a + b

    # Store
    tl.store(C + offs, c, mask=mask)


def triton_vecadd(A, B, C):
    N = A.numel()

    # Fixed optimal config based on your 235us result
    # Tune BLOCK_SIZE based on what worked best in autotuning
    BLOCK_SIZE = 2048  # Start with this, adjust based on your best run
    num_warps = 8

    grid = (triton.cdiv(N, BLOCK_SIZE),)
    vecadd_fp16_kernel[grid](
        A, B, C, N,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
        num_stages=1,
    )
    return C


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


def generate_input(size: int, seed: int) -> input_t:
    """
    Generates random input tensors of specified shapes.
    Returns:
        Tuple of tensors [A, B, C] to be added.
    """
    gen = torch.Generator(device="cuda")
    gen.manual_seed(seed)
    A = torch.randn(
        size, size, device="cuda", dtype=torch.float16, generator=gen
    ).contiguous()
    B = torch.randn(
        size, size, device="cuda", dtype=torch.float16, generator=gen
    ).contiguous()
    C = torch.empty(size, size, device="cuda", dtype=torch.float16).contiguous()
    return A, B, C


def custom_kernel(data: input_t) -> output_t:
    """Fixed optimal Triton config - no autotuning variance"""
    with DeterministicContext():
        A, B, C = data
        return triton_vecadd(A, B, C)


check_implementation = make_match_reference(ref_kernel)