submission 67714

#!POPCORN leaderboard vectoradd_v2

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


@triton.autotune(
    configs=[
        triton.Config({"BLOCK_SIZE": 8192},  num_warps=8, num_stages=6),
        triton.Config({"BLOCK_SIZE": 4096},  num_warps=8, num_stages=5),
        triton.Config({"BLOCK_SIZE": 16384}, num_warps=8, num_stages=6),
    ],
    key=[],
)
@triton.jit
def vecadd_kernel(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    SPAN: tl.constexpr = BLOCK_SIZE * 4  # power-of-two span → OK for arange
    offs = pid * SPAN + tl.arange(0, SPAN)
    mask = offs < N

    # Encourage coalesced 128B vector loads
    tl.multiple_of(offs, 128)
    tl.max_contiguous(offs, SPAN)

    # Fast streaming load (L2 only), use fp16 math pipeline friendly pattern
    a = tl.load(A_ptr + offs, mask=mask, other=0.0, cache_modifier=".cg")
    b = tl.load(B_ptr + offs, mask=mask, other=0.0, cache_modifier=".cg")
    c = a + b

    # Avoid extra hazards — async commit path on Hopper is faster for large N
    tl.store(C_ptr + offs, c, mask=mask)


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

    # ensure contiguous access (esp. for Popcorn harness copies)
    if not A.is_contiguous(): A = A.contiguous()
    if not B.is_contiguous(): B = B.contiguous()

    grid = lambda META: (triton.cdiv(N, META["BLOCK_SIZE"] * 4),)
    vecadd_kernel[grid](A, B, C, N)
    return C