# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import functools

import torch
from torch import Tensor
from torch.library import Library

from vllm import ir
from vllm.platforms import current_platform
from vllm.utils.torch_utils import direct_register_custom_op

current_platform.import_kernels()


def is_aiter_found() -> bool:
    from importlib.util import find_spec

    return find_spec("aiter") is not None


aiter_lib = Library("vllm_aiter", "FRAGMENT")
"""
This library holds torch custom ops for wrapped AITER ops.
Many AITER ops want to remain invisible to torch.compile even after lowering.
They are thus wrapped into torch custom ops inside the IR op implementations.
"""

direct_register_aiter_op = functools.partial(
    direct_register_custom_op, target_lib=aiter_lib
)
"""Syntactic sugar for registering AITER custom ops."""

AITER_SUPPORTED = is_aiter_found()
"""Most kernels in this file are supported if AITER is installed."""


def flatten_to_2d_is_free(x: Tensor) -> bool:
    """Whether ``x.reshape(-1, x.shape[-1])`` is a view with unit-stride rows.

    The AITER norm kernels only take dense 2D inputs, so the wrappers below
    flatten the leading dims with ``Tensor.reshape``, which silently falls back
    to ``contiguous()`` when the flattened shape is not expressible with the
    existing strides, adding a whole-tensor device copy that costs more than
    the kernel saves. A strided last dim flattens without a copy but not into a
    dense buffer, so it is rejected too.
    """
    if x.dim() <= 2:
        return True
    if x.stride(-1) != 1:
        return False
    expected_stride = x.size(-1)
    for i in range(x.dim() - 2, -1, -1):
        if x.size(i) != 1 and x.stride(i) != expected_stride:
            return False
        expected_stride *= x.size(i)
    return True


rms_no_var_16bit_only = (
    lambda x, weight, epsilon, variance_size=None: variance_size is None
    and x.dtype in (torch.float16, torch.bfloat16)
    and (weight is None or weight.dtype == x.dtype)
    and flatten_to_2d_is_free(x)
)
"""AITER rms_norm only supports float16 and bfloat16 acts, no var_size override,
requires weight dtype to match x dtype, and requires flattening the activation
to 2D to be free."""


@ir.ops.rms_norm.register_impl(
    "aiter", supports_args=rms_no_var_16bit_only, supported=AITER_SUPPORTED
)
def rms_norm(
    x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
) -> Tensor:
    assert variance_size is None
    assert x.dtype in (torch.float16, torch.bfloat16)
    if weight is None:
        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
    return torch.ops.vllm_aiter.rms_norm(x, weight, epsilon)


def _rms_norm_impl(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor:
    from aiter import rms_norm

    if x.dim() > 2:
        x_original_shape = x.shape
        x = x.reshape(-1, x_original_shape[-1])
        x = rms_norm(x, weight, variance_epsilon)
        return x.reshape(x_original_shape)

    return rms_norm(x, weight, variance_epsilon)


def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor:
    return torch.empty_like(x)


direct_register_aiter_op(
    op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake
)
rms_add_no_var_16bit_only = (
    lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
    and x.dtype in (torch.float16, torch.bfloat16)
    and (weight is None or weight.dtype == x.dtype)
    and flatten_to_2d_is_free(x)
    and flatten_to_2d_is_free(x_residual)
)
"""
AITER fused_add_rms_norm only supports 16-bit activations and no var_size override.
Requires weight dtype to match x dtype, and flattening both the activation and the
residual to 2D to be free.
"""


@ir.ops.fused_add_rms_norm.register_impl(
    "aiter", supports_args=rms_add_no_var_16bit_only, supported=AITER_SUPPORTED
)
def fused_add_rms_norm(
    x: Tensor,
    x_residual: Tensor,
    weight: Tensor | None,
    epsilon: float,
    variance_size: int | None = None,
) -> tuple[Tensor, Tensor]:
    assert variance_size is None
    assert x.dtype in (torch.float16, torch.bfloat16)
    if weight is None:
        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
    return torch.ops.vllm_aiter.fused_add_rms_norm(x, x_residual, weight, epsilon)


def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
    x: torch.Tensor,
    residual: torch.Tensor,
    weight: torch.Tensor,
    variance_epsilon: float,
    gemma_norm: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    from aiter import rmsnorm2d_fwd_with_add

    x_shape = x.shape
    residual_shape = residual.shape
    hidden_size = x.shape[-1]
    x = x.reshape(-1, hidden_size)
    residual = residual.reshape(-1, hidden_size)

    # TODO can out = x and residual_out = residual to save memory?
    #  Need to check if the kernel supports in-place residual output
    #  (if yes set mutates_args and inplace)
    residual_out = torch.empty_like(residual)
    out = torch.empty_like(x)
    rmsnorm2d_fwd_with_add(
        out,  # output
        x,  # input
        residual,  # residual input
        residual_out,  # residual output
        weight,
        variance_epsilon,
        gemma_norm=gemma_norm,
    )
    return out.reshape(x_shape), residual_out.reshape(residual_shape)


def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
    x: torch.Tensor,
    residual: torch.Tensor,
    weight: torch.Tensor,
    variance_epsilon: float,
    gemma_norm: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    residual_out = torch.empty_like(residual)
    out = torch.empty_like(x)
    return out, residual_out


direct_register_aiter_op(
    op_name="fused_add_rms_norm",
    op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl,
    fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake,
)