import torch
import torch_mlu
import torch_mlu_ops as ops
from common_utils import *
from itertools import product

dtype_dict = {
  torch.half: "half",
  torch.bfloat16: "bfloat16",
}
dtype_list = [torch.half]
if torch_mlu.mlu.is_bf16_supported():
    dtype_list.append(torch.bfloat16)

# M<=INT32_MAX, nk<=INT32_MAX, k>=16, n>=16
class TestSmoothQuantMatmulOp(BtTestCase):
    def op_impl_base(self, *args):
        a, a_scale, b, b_scale, dtype, bias, c, act_mode, alpha, beta, use_hp_active = args
        if a_scale is not None:
            a = torch.mul(a, a_scale.unsqueeze(-1)).to(dtype)
        if b_scale is not None:
            b = torch.mul(b, b_scale.unsqueeze(-1)).to(dtype)
        output = torch.matmul(a, b.permute(1,0))
        if bias is not None:
            output += bias
        output = torch.mul(output, alpha)
        if c is not None:
          residual = torch.mul(c, beta)
          output = torch.add(output, residual)
        if act_mode != "none":
            act = act_mode_dict[act_mode]
            output = act(output)
        return output

    def test_smooth_quant_matmul(self):
        m_list = [32, 64, 128]
        n_list = [64, 128, 256]
        k_list = [128, 256, 512]
        has_bias_list = [True, False]
        has_c_list = [True, False]
        act_mode_list = ["none", "silu", "gelu"]
        use_hp_active_list = [True, False]

        args = product(m_list, n_list, k_list, has_bias_list, has_c_list, act_mode_list, dtype_list, use_hp_active_list)
        for m, n, k, has_bias, has_c, act_mode, dtype, use_hp_active in args:
            if has_c and act_mode != "none":
                continue
            a = torch.randn(m, k, device="mlu", dtype=dtype)
            b = torch.randn(n, k, device="mlu", dtype=dtype)
            bias, c = None, None
            if has_bias:
                bias = torch.randn(n, device="mlu", dtype=dtype)
            if has_c:
                c = torch.randn(m, n, device="mlu", dtype=dtype)

            input_smooth = torch.randn(k, device="mlu", dtype=torch.float).abs() + 0.1
            quant_input, input_scale = QuantByRow(a * input_smooth, 8)
            quant_weight, weight_scale = QuantByRow(b / input_smooth, 8)

            torch_output = self.op_impl_base(quant_input, input_scale, quant_weight, weight_scale, dtype, bias, c,
                                             act_mode, 1.0, 1.0, use_hp_active)
            tmo_output = ops.smooth_quant_matmul(quant_input, input_scale, quant_weight, weight_scale, dtype, bias, c,
                                                 act_mode, 1.0, 1.0, use_hp_active)
            self.assertTensorsEqual(tmo_output.cpu().float(), torch_output.cpu().float(), 0.006, use_MSE=True)

    def test_inductor(self):
        has_bias_list = [True, False]
        has_c_list = [True, False]
        act_mode_list = ["none", "silu", "gelu"]
        arg = product(has_c_list, has_bias_list, act_mode_list, dtype_list)
        for has_c, has_bias, act_mode, dtype in arg:
            if has_c and act_mode != "none":
                continue
            print(f"===has_c: {has_c}, has_bias: {has_bias}, act_mode: {act_mode}, dtype: {dtype}===")
            M, K, N = 2, 16, 32
            quant_bit_size, use_hp_active, act_coef, alpha, beta, trans_a, trans_b = 8, True, 1., 0.8, 0.3, False, True
            a = torch.randint(0, 10, (M, K), dtype=torch.int8).mlu()
            b = torch.randint(0, 10, (N, K), dtype=torch.int8).mlu()
            c = torch.randn(M, N, device="mlu", dtype=dtype) if has_c else None
            bias = torch.randn(N, device="mlu", dtype=dtype) if has_bias else None
            a_scale = torch.randn(M, device="mlu", dtype=torch.float)
            b_scale = torch.randn(N, device="mlu", dtype=torch.float)
            a_zero, b_zero, c_zero = None, None, None
            c_scale, gemm_output_scale, gemm_output_zero = None, None, None
            quant_algo, a_quant_layout, b_quant_layout = "smooth_quant", "quantize_per_token", "quantize_per_channel"
            str_dtype = dtype_dict[dtype]

            args = [a, a_scale, a_zero,
                    b, b_scale, b_zero,
                    bias, c, c_scale, c_zero,
                    gemm_output_scale, gemm_output_zero,
                    str_dtype, None, quant_algo,
                    a_quant_layout, b_quant_layout,
                    quant_bit_size, act_mode, use_hp_active, act_coef,
                    alpha, beta, trans_a, trans_b,]
            self.base_opcheck(torch.ops.torch_mlu_ops.quant_matmul, args)

if __name__ == '__main__':
    random.seed(0)
    torch.manual_seed(0)
    exit(run_unittest(TestSmoothQuantMatmulOp))
