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 TestWeightOnlyQuantMatmulOp(BtTestCase):
    def op_impl_base(self, *args):
        a, b, scale, zero, bias, c, act_mode, quant_bit_size, alpha, beta, use_hp_active = args
        if quant_bit_size == 4:
            n = b.shape[0]
            b = UnpackInt4(b).view(n, -1)
        if scale is not None:
            if scale.dim() == 2:
                group_size = b.size(1) // scale.size(1)
                scale_bd = scale.unsqueeze(-1).repeat(1, 1, group_size).reshape(b.shape)
            else:
                scale_bd = scale.unsqueeze(-1)
            b = torch.mul(b, scale_bd).to(a.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_weight_only_quant_matmul(self):
        m_list = [32, 64, 128]
        n_list = [128, 256, 512]
        group_list = [1, 2]
        k_list = [128, 256, 512]
        quant_bit_list = [8, 4]
        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, group_list, quant_bit_list, has_bias_list, has_c_list, act_mode_list, dtype_list, use_hp_active_list)
        for m, n, k, group, quant_bit, 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
            zero = None
            if has_bias:
                bias = torch.randn(n, device="mlu", dtype=dtype)
            if has_c:
                c = torch.randn(m, n, device="mlu", dtype=dtype)
            quant_weight_int8, weight_scale = QuantByRow(b, quant_bit, group)
            if group != 1:
                if act_mode != "none":
                    continue
                weight_scale = weight_scale.to(a.dtype)
            if quant_bit == 4:
                quant_weight_int4 = PairlyPackInt8(quant_weight_int8)

            args = (a, quant_weight_int4 if quant_bit == 4 else quant_weight_int8,
                        weight_scale, zero, bias, c, act_mode, quant_bit, 1.0, 1.0, use_hp_active)
            torch_output = self.op_impl_base(*args)
            tmo_output = ops.weight_only_quant_matmul(*args)
            self.assertTensorsEqual(tmo_output.cpu().float(), torch_output.cpu().float(), 0.004, use_MSE=True)

    def test_inductor(self):
        M, K, N, group_num = 2, 256, 32, 4
        quant_bit_size, act_mode, use_hp_active, act_coef, alpha, beta, trans_a, trans_b = 8, 'none', True, 1., 0.8, 0.3, False, True
        has_res_list = [True, False]
        group_quant_list = [True, False]
        args = product(has_res_list, group_quant_list, dtype_list)
        for has_res, group_quant, dtype in args:
            print(f"==has_res: {has_res}, group_quant: {group_quant}, dtype: {dtype}==")
            a = torch.randn((M, K), dtype=dtype).mlu()
            b = torch.randint(-128, 127, (N, K), dtype=torch.int8).mlu()
            c = torch.randn(M, N, device="mlu", dtype=dtype) if has_res else None
            bias = torch.randn(N, device="mlu", dtype=dtype)
            group_wise_scale = torch.randn((N, group_num), device="mlu", dtype=dtype)
            b_quant_layout = "quantize_group_wise" if group_quant else "quantize_per_channel"
            b_scale = group_wise_scale if group_quant else None
            gemm_output_scale =  None if group_quant else torch.randn(N, device="mlu", dtype=torch.float)
            a_scale, a_zero, b_zero, c_zero = None, None, None, None
            c_scale, gemm_output_zero = None, None
            quant_algo, a_quant_layout = "weight_only", "quantize_none"
            dtype_str = 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,
                    dtype_str, 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(TestWeightOnlyQuantMatmulOp))
