from functools import reduce
from typing import Any, Tuple

import torch
from torch import Tensor
from torch.cuda.amp import custom_bwd, custom_fwd

try:
    import triton
    import triton.language as tl

    HAS_TRITON = True
except ImportError:
    HAS_TRITON = False
    print("please install triton from https://github.com/openai/triton")

if HAS_TRITON:
    PRECISION_MAP = {
        "fp32": (0, torch.float32),
        "fp16": (1, torch.float16),
        "bf16": (2, torch.bfloat16),
    }

    @triton.jit
    def _llama_act_combine_forward(
        X_GATE1,
        X_GATE2,
        X_UP,
        Y,
        stride,  # how much to increase the pointer when moving by 1 row
        N,  # number of columns in X
        BLOCK_SIZE: tl.constexpr,
    ):
        # Map the program id to the row of X and Y it should compute.
        row = tl.program_id(0)
        X_GATE1 += row * stride
        X_GATE2 += row * stride
        X_UP += row * stride
        Y += row * stride

        # do activation and combine, and store in y
        for off in range(0, N, BLOCK_SIZE):
            cols = off + tl.arange(0, BLOCK_SIZE)
            mask = cols < N
            x_gate1 = tl.load(X_GATE1 + cols, mask=mask, other=0.0)
            x_gate2 = tl.load(X_GATE2 + cols, mask=mask, other=0.0)
            x_up = tl.load(X_UP + cols, mask=mask, other=0.0)
            x_gate2_sigmoid = tl.sigmoid(x_gate2.to(tl.float32)).to(x_gate2.dtype)
            y = x_gate1 * x_gate2 * x_gate2_sigmoid * x_up
            # Write output
            tl.store(Y + cols, y, mask=mask)

    @triton.jit
    def _llama_act_combine_backward(
        X_GATE1,
        X_GATE2,
        X_UP,
        X_GATE1_GRAD,
        X_GATE2_GRAD,
        X_UP_GRAD,
        Y_GRAD,
        stride,  # how much to increase the pointer when moving by 1 row
        N,  # number of columns in X
        BLOCK_SIZE: tl.constexpr,
    ):
        # Map the program id to the row of X and Y it should compute.
        row = tl.program_id(0)
        X_GATE1 += row * stride
        X_GATE2 += row * stride
        X_UP += row * stride
        X_GATE1_GRAD += row * stride
        X_GATE2_GRAD += row * stride
        X_UP_GRAD += row * stride
        Y_GRAD += row * stride

        # do activation and combine, and store in y
        for off in range(0, N, BLOCK_SIZE):
            cols = off + tl.arange(0, BLOCK_SIZE)
            mask = cols < N
            x_gate1 = tl.load(X_GATE1 + cols, mask=mask, other=0.0)
            x_gate2 = tl.load(X_GATE2 + cols, mask=mask, other=0.0)
            x_up = tl.load(X_UP + cols, mask=mask, other=0.0)
            y_grad = tl.load(Y_GRAD + cols, mask=mask, other=0.0)

            # forward: y = x_gate1 * x_gate2 * tl.sigmoid(x_gate2) * x_up
            x_gate2_sigmoid = tl.sigmoid(x_gate2.to(tl.float32)).to(x_gate2.dtype)
            x_gate2_act = y_grad * x_gate2 * x_gate2_sigmoid
            x_up_grad = x_gate2_act * x_gate1
            x_gate1_grad = x_gate2_act * x_up
            # grad(x*sigmoid(x)) = sigmoid(x) + x * sigmoid(x) * [1 − sigmoid(x)]
            #                    = sigmoid(x) * {1 + x * [(1 − sigmoid(x)]}
            x_gate2_grad = (y_grad * x_gate1 * x_up) * x_gate2_sigmoid * (1 + x_gate2 * (1 - x_gate2_sigmoid))

            # Write output
            tl.store(X_GATE1_GRAD + cols, x_gate1_grad, mask=mask)
            tl.store(X_GATE2_GRAD + cols, x_gate2_grad, mask=mask)
            tl.store(X_UP_GRAD + cols, x_up_grad, mask=mask)

    class LlamaActCombine(torch.autograd.Function):
        """
        act(x_gate) * x_up

        Args:
            x_gate (torch.Tensor): (b, l, 2d) x_gate
            x_up (torch.Tensor): (b, l, d) x_up
            activation (str): only support swiglu
            precision (str): fp32, fp16, bf16
        """

        @staticmethod
        @custom_fwd
        def forward(ctx: Any, x_gate: torch.Tensor, x_up: torch.Tensor, activation: str = "swiglu") -> torch.Tensor:
            """
            act(x_gate) * x_up

            Args:
                x_gate (torch.Tensor): (b, l, 2d) x gate
                x_up (torch.Tensor): (b, l, d) x up
                activation (str): only support swiglu
            """
            assert activation == "swiglu", "Only swiglu is supported"

            # split x gate
            assert x_gate.shape[-1] % 2 == 0, "axis size must be divisible by 2"
            x_gate1, x_gate2 = torch.split(x_gate, x_gate.shape[-1] // 2, -1)
            x_gate1 = x_gate1.contiguous()
            x_gate2 = x_gate2.contiguous()
            if not x_up.is_contiguous():
                x_up = x_up.contiguous()
            # assert shape
            assert x_gate1.shape == x_gate2.shape == x_up.shape

            # add ctx for backward
            if x_gate.requires_grad:
                ctx.save_for_backward(x_gate1, x_gate2, x_up)

            # allocate output
            y = torch.empty_like(x_up)
            M, N = reduce(lambda x, y: x * y, x_up.shape[:-1]), x_up.shape[-1]

            # Less than 64KB per feature: enqueue fused kernel
            MAX_FUSED_SIZE = 65536 // x_gate.element_size()
            BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
            if N > BLOCK_SIZE:
                raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
            # heuristics for number of warps
            num_warps = min(max(BLOCK_SIZE // 256, 1), 8)
            # restore setting
            ctx.M, ctx.N, ctx.BLOCK_SIZE, ctx.num_warps = M, N, BLOCK_SIZE, num_warps
            # enqueue kernel
            _llama_act_combine_forward[(M,)](
                x_gate1, x_gate2, x_up, y, x_up.stride(-2), N, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps
            )
            return y

        @staticmethod
        @custom_bwd
        def backward(ctx: Any, *grad_outputs: Tensor) -> Tuple[Tensor, Tensor, None, None]:
            # restore from ctx
            (x_gate1, x_gate2, x_up) = ctx.saved_tensors
            M, N, BLOCK_SIZE, num_warps = ctx.M, ctx.N, ctx.BLOCK_SIZE, ctx.num_warps

            # init grad
            y_grad = grad_outputs[0]
            x_gate1_grad, x_gate2_grad, x_up_grad = (
                torch.empty_like(x_gate1),
                torch.empty_like(x_gate2),
                torch.empty_like(x_up),
            )

            # enqueue kernel
            _llama_act_combine_backward[(M,)](
                x_gate1,
                x_gate2,
                x_up,
                x_gate1_grad,
                x_gate2_grad,
                x_up_grad,
                y_grad,
                x_up.stride(-2),
                N,
                BLOCK_SIZE=BLOCK_SIZE,
                num_warps=num_warps,
            )
            x_gate_grad = torch.cat([x_gate1_grad, x_gate2_grad], dim=-1)
            return x_gate_grad, x_up_grad, None, None