# SPDX-License-Identifier: MIT
 
import torch
from torch import Tensor
from ..jit.core import compile_ops
from typing import Optional

MD_NAME = "module_rmsnorm"


@compile_ops("module_rmsnorm")
def rms_norm_cu(
    out: Tensor,
    input: Tensor,
    weight: Tensor,
    epsilon: float,
) -> None:
    """
    Cuda version of rmsnorm
    """
    ...


@compile_ops("module_rmsnorm")
def fused_add_rms_norm_cu(
    input: Tensor,  # input/out
    residual_in: Tensor,  # residual_in/out
    weight: Tensor,
    epsilon: float,
) -> None:
    """
    Cuda version of rmsnorm fused add
    """
    ...


def gen_rms_norm_fake_tensor(
    input: Tensor,
    weight: Tensor,
    epsilon: float,
) -> Tensor:
    return torch.empty_like(input, dtype=input.dtype, device=input.device)


@compile_ops(
    "module_rmsnorm", fc_name="rmsnorm2d_fwd", gen_fake=gen_rms_norm_fake_tensor
)
def rms_norm(
    input: Tensor,
    weight: Tensor,
    epsilon: float,
) -> Tensor:
    """
    CK version of rmsnorm
    """
    ...


@compile_ops("module_rmsnorm", gen_fake=gen_rms_norm_fake_tensor)
def rmsnorm2d_fwd(
    input: torch.Tensor,
    weight: torch.Tensor,
    epsilon: float,
) -> Tensor: ...


@compile_ops("module_rmsnorm")
def rmsnorm2d_fwd_with_add(
    out: Tensor,
    input: Tensor,
    residual_in: Tensor,
    residual_out: Tensor,
    weight: Tensor,
    epsilon: float,
) -> None: ...


@compile_ops("module_rmsnorm")
def rmsnorm2d_fwd_with_smoothquant(
    out: Tensor,
    input: Tensor,
    xscale: Tensor,
    yscale: Tensor,
    weight: Tensor,
    epsilon: float,
) -> None: ...


@compile_ops("module_rmsnorm")
def rmsnorm2d_fwd_with_add_smoothquant(
    out: Tensor,
    input: Tensor,
    residual_in: Tensor,
    residual_out: Tensor,
    xscale: Tensor,
    yscale: Tensor,
    weight: Tensor,
    epsilon: float,
    out_before_quant: Optional[Tensor] = None,
) -> None: ...


@compile_ops("module_rmsnorm")
def rmsnorm2d_fwd_with_dynamicquant(
    out: Tensor,
    input: Tensor,
    yscale: Tensor,
    weight: Tensor,
    epsilon: float,
) -> None: ...


@compile_ops("module_rmsnorm")
def rmsnorm2d_fwd_with_add_dynamicquant(
    out: Tensor,
    input: Tensor,
    residual_in: Tensor,
    residual_out: Tensor,
    yscale: Tensor,
    weight: Tensor,
    epsilon: float,
) -> None: ...


@compile_ops("module_rmsnorm", gen_fake=gen_rms_norm_fake_tensor)
def head_rms_norm(
    input: Tensor,       # [num_tokens, num_heads * head_dim]
    weight: Tensor,      # [num_heads * head_dim]
    epsilon: float,
    norm_head_dim: int,  # head_dim (size of each head's normalization window)
) -> Tensor:
    """
    Apply RMS normalization per head independently.

    Unlike standard rms_norm which normalizes over the entire last dimension,
    head_rms_norm normalizes each head's head_dim elements separately with
    its own weight parameters.

    Args:
        input: shape [num_tokens, num_heads * head_dim]
        weight: shape [num_heads * head_dim]
        epsilon: small value for numerical stability
        norm_head_dim: the dimension of each head (head_dim)

    Returns:
        Tensor with same shape as input
    """
    ...
