import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *

class TestFusedNormResidualFFNoP(BtTestCase):
    def op_impl_base(self, *args):
        input, up_fc_weight, up_fc_bias, down_proj_weight, down_proj_bias, gate_up_proj_weight, \
            gate_up_proj_bias, layernorm_weight, layernorm_bias, eps, act_mode, residual_is, \
                alpha, beta = args
        hidden_size = input.size(-1)
        layernorm = torch.nn.LayerNorm(hidden_size)
        layernorm.weight = torch.nn.Parameter(layernorm_weight)
        layernorm.bias = torch.nn.Parameter(layernorm_bias)
        layernorm.eps = eps
        act = act_mode_dict[act_mode]

        residual = input
        norm_out = input
        if layernorm_weight is not None:
            norm_out = layernorm(input)
        if residual_is == "normed_input":
            residual = norm_out
        up_fc_out = torch.matmul(norm_out, up_fc_weight.permute(1, 0)) + up_fc_bias
        act_out = act(up_fc_out.float()).to(up_fc_out.dtype)
        if gate_up_proj_weight is not None:
            gate_up_proj_out = torch.matmul(norm_out, gate_up_proj_weight.permute(1, 0)) + gate_up_proj_bias
            down_proj_out = torch.matmul(act_out * gate_up_proj_out, down_proj_weight.permute(1, 0)) + down_proj_bias
        else:
            down_proj_out = torch.matmul(act_out, down_proj_weight.permute(1, 0)) + down_proj_bias
        if residual_is != 'none':
            out = beta * residual + alpha * down_proj_out
        else:
            out = down_proj_out
        return out

    def test_ffn(self):
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        for dtype in dtype_list:
            batch, seq_len, hidden_size, inner_size, alpha, beta = 4, 16, 512, 512, 0.5, 0.3
            eps, act_mode, residual_is = 1e-5, "silu", "input"
            print(f"batch: {batch}, seq_len: {seq_len}, hidden_size: {hidden_size}, inner_size: {inner_size}, alpha: {alpha}, beta: {beta}, \
eps: {eps}, act_mode: {act_mode}, residual_is: {residual_is}, dtype: {dtype} testing...", flush=True)
            input = torch.randn((batch, seq_len, hidden_size), dtype=dtype, device="mlu")
            layernorm_weight = torch.randn(hidden_size, dtype=dtype, device="mlu")
            layernorm_bias = torch.normal(0, 0.1, (hidden_size,), dtype=dtype, device="mlu")
            up_fc_weight = torch.randn((inner_size, hidden_size), dtype=dtype, device="mlu")
            up_fc_bias = torch.normal(0, 0.1, (inner_size,), dtype=dtype, device="mlu")
            gated_fc_weight = torch.randn((inner_size, hidden_size), dtype=dtype, device="mlu")
            gated_fc_bias = torch.normal(0, 0.1, (inner_size,), dtype=dtype, device="mlu")
            down_fc_weight = torch.randn((hidden_size, inner_size), dtype=dtype, device="mlu")
            down_fc_bias = torch.normal(0, 0.1, (hidden_size,), dtype=dtype, device="mlu")
            torch_output = self.op_impl_base(input,
                                             up_fc_weight, up_fc_bias,
                                             down_fc_weight, down_fc_bias,
                                             gated_fc_weight, gated_fc_bias,
                                             layernorm_weight, layernorm_bias,
                                             eps, act_mode, residual_is,
                                             alpha, beta)
            tmo_output = ops.fused_norm_residual_ffn(input,
                                                     up_fc_weight, up_fc_bias,
                                                     down_fc_weight, down_fc_bias,
                                                     gated_fc_weight, gated_fc_bias,
                                                     layernorm_weight, layernorm_bias,
                                                     eps, act_mode, residual_is,
                                                     alpha, beta)
            self.assertTensorsEqual(torch_output.cpu().float(), tmo_output.cpu().float(),
                                    0.005, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        batch, seq_len, hidden_size, inner_size, act_mode, dtype = 1, 10, 128, 1024, 'silu', torch.half
        input = torch.randn((batch, seq_len, hidden_size), dtype=dtype, device="mlu")
        layernorm_weight = torch.randn((hidden_size,), dtype=dtype, device="mlu")
        layernorm_bias = torch.randn((hidden_size,), dtype=dtype, device="mlu")
        up_fc_weight = torch.randn((inner_size, hidden_size), dtype=dtype, device="mlu")
        gate_up_proj_weight = torch.randn((inner_size, hidden_size), dtype=dtype, device="mlu")
        down_proj_weight = torch.randn((hidden_size, inner_size), dtype=dtype, device="mlu")
        up_fc_bias = torch.randn((inner_size,), dtype=dtype, device="mlu")
        gate_up_proj_bias = torch.randn((inner_size,), dtype=dtype, device="mlu")
        down_proj_bias = torch.randn((hidden_size,), dtype=dtype, device="mlu")
        args = (input, up_fc_weight, up_fc_bias, down_proj_weight, down_proj_bias,
                gate_up_proj_weight, gate_up_proj_bias, layernorm_weight, layernorm_bias, act_mode, "normed_input", 1e-5, 1., 0.)
        self.base_opcheck(torch.ops.torch_mlu_ops.ffn, args)

if __name__ == '__main__':
    exit(run_unittest(TestFusedNormResidualFFNoP))
