import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from itertools import product
import numpy as np

class TestBatchMatMulOp(BtTestCase):
    def op_impl_base(self, *args):
        a, b, c, alpha, beta, a_scale, b_scale, trans_a, trans_b = args
        if trans_a:
            a = a.transpose(1, 2)
        if trans_b:
            b = b.transpose(1, 2)
        if c is None:
            output_dtype = a.dtype
        else:
            output_dtype = c.dtype
        if a_scale is not None:
            a = torch.div(a, a_scale).to(output_dtype)
            b = torch.div(b, b_scale).to(output_dtype)
        output = alpha * torch.bmm(a, b)
        if c is not None:
            output += beta * c
            c.copy_(output)
        return output

    def test_batch_matmul(self):
        batch_list = [5]
        mat_m_list = [32]
        mat_n_list = [256]
        mat_k_list = [128]
        has_res_list = [False, True]
        trans_a_list = [False, True]
        trans_b_list = [False, True]

        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        alpha = 0.625
        beta = 1.0
        args = product(batch_list, mat_m_list, mat_n_list, mat_k_list, has_res_list, dtype_list, trans_a_list, trans_b_list)

        for batch, mat_m, mat_n, mat_k, has_res, dtype, trans_a, trans_b  in args:
            print("batch={}, m={}, n={}, k={}, has_res={}, dtype={}, trans_a={}, trans_b={} testing...".format(
               batch, mat_m, mat_n, mat_k, has_res, dtype, trans_a, trans_b), flush=True)
            shape_a, shape_b = (batch, mat_m, mat_k), (batch, mat_k, mat_n)
            if trans_a:
                shape_a = (batch, mat_k, mat_m)
            if trans_b:
                shape_b = (batch, mat_n, mat_k)
            input = torch.randn(shape_a, dtype=dtype, device='mlu')
            weight = torch.randn(shape_b, dtype=dtype, device='mlu')
            input8, a_scale = QuantByTensor(input, 8)
            weight8, b_scale = QuantByTensor(weight, 8)
            residual = torch.randn((batch, mat_m, mat_n), dtype=dtype, device='mlu')
            res_bak = residual.clone()
            tmo_output_int8 = torch.zeros((batch, mat_m, mat_n), dtype=dtype, device='mlu')
            torch_output_int8 = tmo_output_int8.clone()
            output = self.op_impl_base(input, weight,
                        residual if has_res else None, alpha, beta, 1.0, 1.0, trans_a, trans_b)
            tmo_output = ops.batch_matmul(input, weight,
                        res_bak if has_res else None, alpha, beta, 1.0, 1.0, trans_a, trans_b)

            self.assertTensorsEqual(tmo_output.cpu().float(), output.cpu().float(),
                                0.004, use_MSE=True, use_RAE=True)
            if dtype != torch.bfloat16:
                self.op_impl_base(input8, weight8, torch_output_int8, alpha, beta, a_scale.item(), b_scale.item(), trans_a, trans_b)
                ops.batch_matmul(input8, weight8, tmo_output_int8, alpha, beta, a_scale.item(), b_scale.item(), trans_a, trans_b)
                self.assertTensorsEqual(tmo_output_int8.cpu().float(), torch_output_int8.cpu().float(),
                                0.004, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        batch, mat_m, mat_n, mat_k, alpha, beta, dtype = 6, 64, 256, 128, 0.8, 0.3, torch.float16
        a = torch.randn((batch, mat_m, mat_k), dtype=dtype, device='mlu')
        b = torch.randn((batch, mat_n, mat_k), dtype=dtype, device='mlu')
        c = torch.randn((batch, mat_m, mat_n), dtype=dtype, device='mlu')
        args = (a, b, c, alpha, beta, 1.0, 1.0, False, True)
        self.base_opcheck(torch.ops.torch_mlu_ops.batch_matmul, args)

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