import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from itertools import product
import torch.multiprocessing as mp
import sys
import os

work_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(work_dir))
from common_utils import *

def mm_split_k(rank, *args):
    torch.manual_seed(0)
    mat_m, mat_n, mat_k, has_bias, has_res, dtype, world_size, block_m = args
    assert mat_k % world_size == 0, f"mat_k{mat_k} must be divisible by tp{world_size}"
    block_k = mat_k // world_size
    start_k = rank * block_k
    end_k = start_k + block_k

    setup(rank, world_size)
    pg = get_default_group()
    cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)

    alpha = 0.625
    beta = 1.0 if has_res else 0.
    input = torch.randn((mat_m, mat_k), dtype=dtype, device='mlu')
    residual = torch.randn((mat_m, mat_n), dtype=dtype, device='mlu') if has_res else None
    weight = torch.randn((mat_n, mat_k), dtype=dtype, device="mlu")
    bias = torch.randn(mat_n, dtype=dtype, device="mlu") if has_bias else None
    pt_output = torch.matmul(input, weight.permute(1, 0))
    if has_bias:
        pt_output = pt_output + bias
    pt_output *= alpha
    if has_res:
        pt_output += beta * residual
    if has_bias:
        bias *= alpha
    input = input[..., start_k:end_k].contiguous()
    weight = weight[..., start_k:end_k].contiguous()
    bias = bias if (bias is not None and rank == 0) else None
    residual = residual if (residual is not None and rank == 0) else None
    beta = beta if (residual is not None and rank == 0) else 0.
    output = ops.matmul_allreduce(cncl_comm, input, weight, bias, residual, alpha, beta, block_m)
    if rank == 0:
        assertTensorsEqual(output.cpu().float(), pt_output.cpu().float(), 0.006, use_MSE=True, use_RAE=True)
    cleanup()


class TestMatMulAllReduceOp(BtTestCase):
    def op_impl_base(self, *args):
        return super().op_impl_base(*args)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_matmul_allreduce due to ASan issues")
    def test_matmul_allreduce(self):
        device_n = torch_mlu.mlu.device_count()
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)

        mat_m_list = [32]
        mat_n_list = [256]
        mat_k_list = [1680]
        has_res_list = [False, True]
        has_bias_list = [False, True]
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        block_m = 4
        args = product(mat_m_list, mat_n_list, mat_k_list, has_bias_list, has_res_list, dtype_list)

        for  mat_m, mat_n, mat_k, has_bias, has_res, dtype in args:
            print("m={}, n={}, k={}, has_bias={}, has_res={}, dtype={}, tp={}, block_m={}, testing...".format(
                mat_m, mat_n, mat_k, has_bias, has_res, dtype, device_n, block_m), flush=True)
            param = [mat_m, mat_n, mat_k, has_bias, has_res, dtype, device_n, block_m]
            mp.spawn(mm_split_k, param, nprocs=device_n, join=True)

    @unittest.skip("not test")
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_matmul_allreduce_random due to ASan issues")
    def test_matmul_allreduce_random(self):
        import random
        random.seed(0)
        device_n = torch_mlu.mlu.device_count()
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        for i in range(10):
            tp_num = random.randint(1, device_n)
            mat_m = random.randint(1, 4096)
            mat_n = random.choice([512, 1024, 2048, 4096])
            k_start = (1024 // tp_num) * tp_num
            mat_k = random.randrange(k_start, 10240, tp_num)
            has_res = random.choice([False, True])
            has_bias = random.choice([False, True])
            dtype = random.choice(dtype_list)
            block_m = random.randint(1, 10)
            assert mat_k % tp_num == 0, f"mat_k{mat_k} must be divisible by tp_num{tp_num}"
            print("m={}, n={}, k={}, has_bias={}, has_res={}, dtype={}, tp_num={}, block_m={}, testing...".format(
                mat_m, mat_n, mat_k, has_bias, has_res, dtype, tp_num, block_m), flush=True)
            param = [mat_m, mat_n, mat_k, has_bias, has_res, dtype, tp_num, block_m]
            mp.spawn(mm_split_k, param, nprocs=tp_num, join=True)

    def test_inductor(self):
        return super().test_inductor()

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