"""Autogenerated Triton kernel for vector addition."""

from __future__ import annotations

import torch
import warnings

try:
    import triton
    import triton.language as tl
except ImportError:  # pragma: no cover - optional dependency
    triton = None
    tl = None

DEFAULT_BLOCK_SIZE = 256


def _ceil_div(x: int, y: int) -> int:
    return (x + y - 1) // y


if triton is not None:

    @triton.jit
    def vector_add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
        pid = tl.program_id(axis=0)
        start = pid * BLOCK_SIZE
        offsets = start + tl.arange(0, BLOCK_SIZE)
        mask = offsets < n_elements
        x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
        y = tl.load(y_ptr + offsets, mask=mask, other=0.0)
        tl.store(output_ptr + offsets, x + y, mask=mask)


def run_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Compute ``x + y`` via Triton when available else Torch fallback."""

    if x.shape != y.shape:
        raise ValueError("Inputs must share the same shape")
    if x.dim() != 1:
        raise ValueError("Only 1D tensors are supported")
    x = x.to(dtype=torch.float32)
    y = y.to(dtype=torch.float32)
    if x.device != y.device:
        raise ValueError("Inputs must reside on the same device")

    fallback = x + y
    if triton is None or not torch.cuda.is_available():
        return fallback
    if x.device.type != "cuda":
        return fallback

    n_elements = x.numel()
    output = torch.empty_like(x)
    block_size = DEFAULT_BLOCK_SIZE
    grid = (_ceil_div(n_elements, block_size),)
    try:
        vector_add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=block_size)
    except Exception as exc:  # pragma: no cover - triton runtime issues
        warnings.warn(
            f"Triton vector_add kernel failed; falling back to torch: {exc}",
            RuntimeWarning,
        )
        return fallback
    return output
