# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
#   https://github.com/fla-org/flash-linear-attention/graphs/contributors

import torch
import torch.nn as nn
import torch.nn.functional as F
import triton
import triton.language as tl

from fla.modules.backends import dispatch
from fla.ops.utils.op import exp, log
from fla.utils import IS_AMD, input_guard

# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576
# https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19
# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
# The optimal maximum block size depends on your hardware, your kernel, and your dtype
MAX_FUSED_SIZE = 65536 // 2
STATIC_WARPS = 32 if not IS_AMD else 16


@triton.jit
def kl_div_kernel(
    logits,
    target_logits,
    loss,
    s_logits,
    s_loss,
    reduction: tl.constexpr,
    N: tl.constexpr,
    V: tl.constexpr,
    BV: tl.constexpr,
):
    # https://github.com/triton-lang/triton/issues/1058
    # If N*V is too large, i_n * stride will overflow out of int32, so we convert to int64
    i_n = tl.program_id(0).to(tl.int64)

    logits += i_n * s_logits
    target_logits += i_n * s_logits

    # m is the max value. use the notation from the paper
    sm = float('-inf')
    tm = float('-inf')
    # d is the sum. use the notation from the paper
    sd, td = 0.0, 0.0

    NV = tl.cdiv(V, BV)
    for iv in range(0, NV):
        o_x = iv * BV + tl.arange(0, BV)
        # for student
        b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf'))
        b_sm = tl.max(b_sl)
        m_new = tl.maximum(sm, b_sm)
        sd = sd * exp(sm - m_new) + tl.sum(exp(b_sl - m_new))
        sm = m_new
        # for teacher
        b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf'))
        b_tm = tl.max(b_tl)
        m_new = tl.maximum(tm, b_tm)
        td = td * exp(tm - m_new) + tl.sum(exp(b_tl - m_new))
        tm = m_new

    b_loss = 0.
    # KL(y_true || y) = exp(y_true) * (log(y_true) - log(y))
    for iv in range(0, NV):
        o_x = iv * BV + tl.arange(0, BV)
        b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf'))
        b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf'))
        b_sp_log = b_sl - sm - log(sd)
        b_tp_log = b_tl - tm - log(td)
        b_sp = exp(b_sp_log)
        b_tp = exp(b_tp_log)
        b_kl = tl.where(o_x < V, b_tp * (b_tp_log - b_sp_log), 0)
        b_dl = -b_tp + b_sp
        b_loss += tl.sum(b_kl)
        if reduction == 'batchmean':
            b_dl = b_dl / N
        tl.store(logits + o_x, b_dl, mask=o_x < V)

    # Normalize the loss by the number of elements if reduction is 'batchmean'
    if reduction == 'batchmean':
        b_loss = b_loss / N

    tl.store(loss + i_n * s_loss, b_loss)


@triton.jit
def elementwise_mul_kernel(
    x,
    g,
    N: tl.constexpr,
    B: tl.constexpr,
):
    """
    This function multiplies each element of the tensor pointed by x with the value pointed by g.
    The multiplication is performed in-place on the tensor pointed by x.

    Parameters:
    x:
        Pointer to the input tensor.
    g:
        Pointer to the gradient output value.
    N (int):
        The number of columns in the input tensor.
    B (int):
        The block size for Triton operations.
    """

    # Get the program ID and convert it to int64 to avoid overflow
    i_x = tl.program_id(0).to(tl.int64)
    o_x = i_x * B + tl.arange(0, B)

    # Load the gradient output value
    b_g = tl.load(g)
    b_x = tl.load(x + o_x, mask=o_x < N)
    tl.store(x + o_x, b_x * b_g, mask=o_x < N)


@dispatch('modules')
def fused_kl_div_forward(
    x: torch.Tensor,
    target_x: torch.Tensor,
    weight: torch.Tensor,
    target_weight: torch.Tensor,
    reduction: str = 'batchmean',
    accumulate_grad_in_fp32: bool = True,
):
    device = x.device

    # ideally, we would like to achieve the same memory consumption as [N, H],
    # so the expected chunk size should be:
    # NC = ceil(V / H)
    # C = ceil(N / NC)
    # for ex: N = 4096*4, V = 32000, H = 4096 ==> NC = 8, C = ceil(N / NC) = 2048
    N, H, V = *x.shape, weight.shape[0]
    BV = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
    # TODO: in real cases, we may need to limit the number of chunks NC to
    # ensure the precisions of accumulated gradients
    NC = min(8, triton.cdiv(V, H))
    C = triton.next_power_of_2(triton.cdiv(N, NC))
    NC = triton.cdiv(N, C)

    grad_dtype = torch.float32 if accumulate_grad_in_fp32 else weight.dtype

    dx = torch.zeros_like(x, device=device)
    dw = torch.zeros_like(weight, device=device, dtype=grad_dtype) if weight is not None else None
    # we use fp32 for loss accumulator
    loss = torch.zeros(N, dtype=torch.float32, device=device)

    for ic in range(NC):
        start, end = ic * C, min((ic + 1) * C, N)
        # [C, N]
        c_sx = x[start:end]
        c_tx = target_x[start:end]
        # when doing matmul, use the original precision
        # [C, V]
        c_sl = F.linear(c_sx, weight)
        c_tl = F.linear(c_tx, target_weight)
        if weight is not None and c_sx.dtype != grad_dtype:
            c_sx = c_sx.to(dtype=grad_dtype)

        # unreduced loss
        c_loss = loss[start:end]

        # Here we calculate the gradient of c_sx in place so we can save memory.
        kl_div_kernel[(c_sx.shape[0],)](
            logits=c_sl,
            target_logits=c_tl,
            loss=c_loss,
            s_logits=c_sl.stride(-2),
            s_loss=c_loss.stride(-1),
            reduction=reduction,
            N=N,
            V=V,
            BV=BV,
            num_warps=STATIC_WARPS,
        )

        # gradient of logits is computed in-place by the above triton kernel and is of shape: C x V
        # thus dx[start: end] should be of shape: C x H
        # additionally, since we are chunking the inputs, observe that the loss and gradients are calculated only
        # on `n_non_ignore` tokens. However, the gradient of the input should be calculated for all tokens.
        # Thus, we need an additional scaling factor of (n_non_ignore/total) to scale the gradients.
        # [C, H]

        dx[start:end] = torch.mm(c_sl, weight)

        if weight is not None:
            torch.addmm(
                input=dw,
                mat1=c_sl.t().to(dtype=grad_dtype),
                mat2=c_sx,
                out=dw,
            )

    loss = loss.sum()
    if dw is not None:
        dw = dw.to(weight)
    return loss, dx, dw


@dispatch('modules')
def fused_kl_div_backward(
    do: torch.Tensor,
    dx: torch.Tensor,
    dw: torch.Tensor,
):
    # If cross entropy is the last layer, do is 1.0. Skip the mul to save time
    if torch.ne(do, torch.tensor(1.0, device=do.device)):
        # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
        # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
        N, H = dx.shape
        B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))

        elementwise_mul_kernel[(triton.cdiv(N * H, B),)](
            x=dx,
            g=do,
            N=N*H,
            B=B,
            num_warps=STATIC_WARPS,
        )

        # handle dw
        if dw is not None:
            V, H = dw.shape
            elementwise_mul_kernel[(triton.cdiv(V * H, B),)](
                x=dw,
                g=do,
                N=V*H,
                B=B,
                num_warps=STATIC_WARPS,
            )

    return dx, dw


class FusedKLDivLossFunction(torch.autograd.Function):

    @staticmethod
    @input_guard
    def forward(
        ctx,
        x: torch.Tensor,
        target_x: torch.Tensor,
        weight: torch.Tensor,
        target_weight: torch.Tensor,
        reduction: str,
        accumulate_grad_in_fp32: bool,
    ):
        loss, dx, dw = fused_kl_div_forward(
            x=x,
            target_x=target_x,
            weight=weight,
            target_weight=target_weight,
            reduction=reduction,
            accumulate_grad_in_fp32=accumulate_grad_in_fp32,
        )
        ctx.save_for_backward(dx, dw)
        return loss

    @staticmethod
    @input_guard
    def backward(ctx, do):
        dx, dw = ctx.saved_tensors
        dx, dw = fused_kl_div_backward(do, dx, dw)
        return dx, None, dw, None, None, None


def fused_kl_div_loss(
    x: torch.Tensor,
    target_x: torch.Tensor,
    weight: torch.Tensor,
    target_weight: torch.Tensor,
    reduction: str = 'batchmean',
    accumulate_grad_in_fp32: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Args:
        x (`torch.Tensor`):
            Tensor of shape `[batch_size * seq_len, hidden_size]`.
        target_x (`torch.Tensor`):
            Frozen teacher input tensor of shape `[batch_size * seq_len, hidden_size]`.
            Must not require gradients.
        weight (`torch.Tensor`):
            Tensor of shape `[vocab_size, hidden_size]`.
        target_weight (`torch.Tensor`):
            Frozen teacher weight tensor of shape `[vocab_size, hidden_size]`.
            Must not require gradients.
        reduction (`str`):
            Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'.
        accumulate_grad_in_fp32 (`bool`):
            Whether to accumulate the student weight gradient in fp32 before casting it back
            to `weight.dtype`. Default: `True`.
    Returns:
        loss
    """
    if target_x.requires_grad or target_weight.requires_grad:
        raise RuntimeError(
            "FusedKLDivLoss treats target_x and target_weight as a frozen teacher and does not compute "
            "gradients for them. Detach target_x/target_weight before calling FusedKLDivLoss, or use "
            "torch.nn.functional.kl_div if teacher gradients are required."
        )
    return FusedKLDivLossFunction.apply(
        x,
        target_x,
        weight,
        target_weight,
        reduction,
        accumulate_grad_in_fp32,
    )


class FusedKLDivLoss(nn.Module):

    def __init__(
        self,
        reduction: str = 'batchmean',
        accumulate_grad_in_fp32: bool = True,
    ):
        """
        Args:
            reduction (`str`):
                Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'.
            accumulate_grad_in_fp32 (`bool`):
                Whether to accumulate the student weight gradient in fp32 before casting it back
                to `weight.dtype`. Default: `True`.
        Note:
            FusedKLDivLoss only computes gradients for `x` and `weight`; `target_x` and
            `target_weight` are treated as frozen teacher tensors and must not require gradients.
        """
        super().__init__()

        assert reduction in ['batchmean'], f"reduction: {reduction} is not supported"

        self.reduction = reduction
        self.accumulate_grad_in_fp32 = accumulate_grad_in_fp32

    def forward(
        self,
        x: torch.Tensor,
        target_x: torch.Tensor,
        weight: torch.Tensor,
        target_weight: torch.Tensor,
    ):
        """
        Args:
            x (`torch.Tensor`):
                Tensor of shape `[batch_size * seq_len, hidden_size]`.
            target_x (`torch.Tensor`):
                Frozen teacher input tensor of shape `[batch_size * seq_len, hidden_size]`.
                Must not require gradients.
            weight (`torch.Tensor`):
                Tensor of shape `[vocab_size, hidden_size]`.
            target_weight (`torch.Tensor`):
                Frozen teacher weight tensor of shape `[vocab_size, hidden_size]`.
                Must not require gradients.
        Returns:
            loss
        """
        loss = fused_kl_div_loss(
            x=x,
            target_x=target_x,
            weight=weight,
            target_weight=target_weight,
            reduction=self.reduction,
            accumulate_grad_in_fp32=self.accumulate_grad_in_fp32,
        )
        return loss