submission 67265

from __future__ import annotations

from typing import Tuple

import torch

from task import input_t, output_t


def _validate_tensors(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> None:
    if not (A.is_cuda and B.is_cuda and C.is_cuda):
        raise RuntimeError("All tensors must reside on CUDA devices.")
    if A.dtype != torch.float16 or B.dtype != torch.float16 or C.dtype != torch.float16:
        raise RuntimeError("All tensors must use torch.float16 dtype.")
    if A.shape != B.shape or A.shape != C.shape:
        raise RuntimeError("Input and output tensors must share identical shapes.")
    if not (A.is_contiguous() and B.is_contiguous() and C.is_contiguous()):
        raise RuntimeError("Input and output tensors must be contiguous.")


def _launch_vector_add(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> None:
    # `torch.add` issues an efficient pointwise kernel and honors the provided output buffer.
    torch.add(A, B, out=C)


def custom_kernel(data: input_t) -> output_t:
    A, B, C = data  # type: ignore[misc]
    _validate_tensors(A, B, C)
    _launch_vector_add(A, B, C)
    return C