submission 510348

from __future__ import annotations

import torch

from task import input_t, output_t

try:
    import triton
    import triton.language as tl

    _HAS_TRITON = True
except Exception:
    triton = None
    tl = None
    _HAS_TRITON = False

_graph_pool = torch.cuda.graph_pool_handle()
_graph_key: tuple[int, int, int, int, int, int] | None = None
_graph: torch.cuda.CUDAGraph | None = None


if _HAS_TRITON:
    @triton.autotune(
        configs=[
            triton.Config({"BLOCK": 512}, num_warps=4, num_stages=2),
            triton.Config({"BLOCK": 1024}, num_warps=4, num_stages=2),
            triton.Config({"BLOCK": 1024}, num_warps=8, num_stages=3),
            triton.Config({"BLOCK": 2048}, num_warps=4, num_stages=3),
            triton.Config({"BLOCK": 2048}, num_warps=8, num_stages=4),
            triton.Config({"BLOCK": 4096}, num_warps=4, num_stages=3),
            triton.Config({"BLOCK": 4096}, num_warps=8, num_stages=4),
            triton.Config({"BLOCK": 8192}, num_warps=4, num_stages=4),
            triton.Config({"BLOCK": 8192}, num_warps=8, num_stages=4),
            triton.Config({"BLOCK": 16384}, num_warps=4, num_stages=5),
            triton.Config({"BLOCK": 16384}, num_warps=8, num_stages=5),
        ],
        key=["n_elements"],
    )
    @triton.jit
    def _vadd_1d(a_ptr, b_ptr, out_ptr, n_elements, BLOCK: tl.constexpr):
        pid = tl.program_id(axis=0)
        offsets = pid * BLOCK + tl.arange(0, BLOCK)
        mask = offsets < n_elements
        tl.multiple_of(offsets, 16)
        tl.max_contiguous(offsets, 256)
        a = tl.load(a_ptr + offsets, mask=mask, other=0.0, cache_modifier=".cg")
        b = tl.load(b_ptr + offsets, mask=mask, other=0.0, cache_modifier=".cg")
        tl.store(out_ptr + offsets, a + b, mask=mask)


def custom_kernel(data: input_t) -> output_t:
    """
    pmpp_v2 vectoradd schema:
      data = (A, B, output)
    Writes A+B into `output` (no temporary allocation) and returns `output`.
    """
    A, B, output = data

    # Assume task-provided inputs are well-formed; keep overhead minimal.
    if (
        A.is_cuda
        and B.is_cuda
        and output.is_cuda
        and A.dtype == torch.float16
        and B.dtype == torch.float16
        and output.dtype == torch.float16
        and A.is_contiguous()
        and B.is_contiguous()
        and output.is_contiguous()
        and A.shape == B.shape
        and A.shape == output.shape
        and A.ndim == 2
        and int(A.shape[0]) <= 2048
    ):
        global _graph_key, _graph
        device = A.device.index or 0
        key = (device, int(A.shape[0]), int(A.shape[1]), A.data_ptr(), B.data_ptr(), output.data_ptr())
        if key == _graph_key and _graph is not None:
            _graph.replay()
            return output

        graph = torch.cuda.CUDAGraph()
        torch.add(A, B, out=output)
        torch.cuda.synchronize()

        with torch.cuda.graph(graph, pool=_graph_pool):
            torch.add(A, B, out=output)
        torch.cuda.synchronize()

        _graph_key = key
        _graph = graph
        return output

    if (
        _HAS_TRITON
        and A.is_cuda
        and B.is_cuda
        and output.is_cuda
        and A.dtype == torch.float16
        and B.dtype == torch.float16
        and output.dtype == torch.float16
        and A.is_contiguous()
        and B.is_contiguous()
        and output.is_contiguous()
        and A.shape == B.shape
        and A.shape == output.shape
        and A.ndim == 2
        and int(A.shape[0]) >= 16384
    ):
        n_elements = A.numel()
        grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK"]),)
        _vadd_1d[grid](A, B, output, n_elements=n_elements)
        return output

    torch.add(A, B, out=output)
    return output