submission 512773

# submission_vectoradd_v2.py
# Element-wise float16 matrix addition: C = A + B
# Interface: custom_kernel((A, B, C)) -> C
#   A, B: (size, size) float16 on CUDA
#   C:    (size, size) float16 pre-allocated output
#
# Strategy: flatten the 2D matrices to 1D, launch a 1D autotuned Triton
# kernel.  Pure memory-bandwidth-bound — goal is to saturate HBM.

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


@triton.autotune(
    configs=[
        triton.Config({'BLOCK': 512},  num_warps=4),
        triton.Config({'BLOCK': 1024}, num_warps=4),
        triton.Config({'BLOCK': 2048}, num_warps=8),
        triton.Config({'BLOCK': 4096}, num_warps=8),
        triton.Config({'BLOCK': 8192}, num_warps=16),
    ],
    key=['N'],
)
@triton.jit
def _vadd_kernel(a_ptr, b_ptr, c_ptr, N: int, BLOCK: tl.constexpr):
    pid  = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N
    a = tl.load(a_ptr + offs, mask=mask)
    b = tl.load(b_ptr + offs, mask=mask)
    tl.store(c_ptr + offs, a + b, mask=mask)


def custom_kernel(data: input_t) -> output_t:
    A, B, C = data
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK']),)
    _vadd_kernel[grid](A, B, C, N)
    return C