import torch
import torch_mlu
import unittest
import math
import torch_mlu_ops as ops
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 flash_attn_sq_mm(q, k, v, smooth, quant_weight, weight_scale,
                     bias, softmax_scale, is_causal, world_size = 1):
    q_list = q.chunk(world_size, dim=2)
    k_list = k.chunk(world_size, dim=2)
    v_list = v.chunk(world_size, dim=2)
    smooth_list = smooth.chunk(world_size, dim=0)
    quant_weight_list = quant_weight.chunk(world_size, dim=1)
    quant_weight_list = [w.contiguous() for w in quant_weight_list]
    output1 = torch.zeros(q.size(0) * q.size(1), q.size(2) * q.size(3), dtype=q.dtype).mlu()
    for i in range(world_size):
        attn_output = ops.flash_attention(q_list[i], k_list[i], v_list[i], None, None, None,
            None, None, q.size(1), k.size(1), softmax_scale, is_causal).flatten(-2, -1).flatten(0, 1)
        quant_input, input_scale = ops.per_token_smooth_quantize(attn_output, smooth_list[i], None)
        output1 += ops.smooth_quant_matmul(quant_input, input_scale,
            quant_weight_list[i], weight_scale, q.dtype, bias if i == 0 else None)
    attn_output = ops.flash_attention(q, k, v, None, None, None,
        None, None, q.size(1), k.size(1), softmax_scale, is_causal).flatten(-2, -1).flatten(0, 1)
    quant_input, input_scale = ops.per_token_smooth_quantize(attn_output, smooth, None)
    output2 = ops.smooth_quant_matmul(quant_input, input_scale, quant_weight, weight_scale, q.dtype, bias)
    return output1, output2

def tp_flash_attn_sq_mm(rank, *args):
    world_size, q, k, v, smooth, quant_weight, weight_scale, bias, softmax_scale, is_causal, base_output = args
    q_cpu, k_cpu, v_cpu, smooth_cpu = q.cpu(), k.cpu(), v.cpu(), smooth.cpu()
    quant_weight_cpu, weight_scale_cpu, bias_cpu = quant_weight.cpu(), weight_scale.cpu(), bias.cpu()
    base_output_cpu = base_output.cpu()
    setup(rank, world_size)
    pg = get_default_group()
    cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
    head_num_q = q.size(2)
    head_num_kv = k.size(2)
    seq = q.size(1)
    assert head_num_q % world_size == 0
    assert head_num_kv % world_size == 0
    head_num_q_tp = head_num_q // world_size
    head_num_kv_tp = head_num_kv // world_size
    q = q_cpu.mlu()
    k = k_cpu.mlu()
    v = v_cpu.mlu()
    smooth = smooth_cpu.mlu()
    quant_weight = quant_weight_cpu.mlu()
    weight_scale = weight_scale_cpu.mlu()
    # Note: only tp0 add bias
    bias = bias_cpu.mlu() if rank == 0 else None
    q_list = q.chunk(world_size, dim=2)
    k_list = k.chunk(world_size, dim=2)
    v_list = v.chunk(world_size, dim=2)
    smooth_list = smooth.chunk(world_size, dim=0)
    quant_weight_list = quant_weight.chunk(world_size, dim=1)
    quant_weight_list = [w.contiguous() for w in quant_weight_list]
    # test pad mode
    output_pad = ops.flash_attn_sq_mm_allreduce(cncl_comm, q_list[rank], k_list[rank], v_list[rank],
        None, None, None, None,
        smooth_list[rank], quant_weight_list[rank], weight_scale, bias, seq,
        seq, softmax_scale, is_causal)
    assertTensorsEqual(output_pad.cpu().float(), base_output_cpu.float(), 0.006, use_MSE=True, use_RAE=True)
    # test pack mode
    cu_seq_lens_q = torch.tensor([0, seq], dtype=torch.int32).mlu()
    cu_seq_lens_k = torch.tensor([0, seq], dtype=torch.int32).mlu()
    q_pack = q_list[rank].flatten(0, 1)
    k_pack = k_list[rank].flatten(0, 1)
    v_pack = v_list[rank].flatten(0, 1)
    output_pack = ops.flash_attn_sq_mm_allreduce(cncl_comm, q_pack, k_pack, v_pack,
        cu_seq_lens_q, cu_seq_lens_k, None, None,
        smooth_list[rank], quant_weight_list[rank], weight_scale, bias, seq,
        seq, softmax_scale, is_causal)
    assertTensorsEqual(output_pack.cpu().float(), base_output_cpu.float(), 0.006, use_MSE=True, use_RAE=True)
    cleanup()

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

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_flash_attn_sq_mm_split_seq due to ASan issues")
    def test_flash_attn_sq_mm_split_seq(self):
        batch, seq, head_num_q, head_num_kv, head_size, is_causal, block_seq = 16, 1024, 8, 1, 128, True, 4
        dtype = torch.bfloat16 if torch_mlu.mlu.is_bf16_supported() else torch.half
        hidden_size = head_num_q * head_size
        softmax_scale = 1 / math.sqrt(head_size)
        qkv = torch.randn(batch, seq, head_num_q + 2 * head_num_kv, head_size, dtype=dtype).mlu()
        smooth = torch.zeros(hidden_size, dtype=torch.float).mlu() + 1.0
        bias = torch.zeros(hidden_size, dtype=dtype).mlu()
        weight = torch.randn(hidden_size, hidden_size, dtype=dtype).mlu()
        quant_weight, weight_scale = QuantByRow(weight / smooth, 8)
        q = qkv[:, :, :head_num_q, :]
        k = qkv[:, :, head_num_q:head_num_q+head_num_kv, :]
        v = qkv[:, :, head_num_q+head_num_kv:, :]
        attn_output = ops.flash_attention(q, k, v, None, None, None,
            None, None, q.size(1), k.size(1), softmax_scale, is_causal).flatten(-2, -1).flatten(0, 1)
        quant_input, input_scale = ops.per_token_smooth_quantize(attn_output, smooth, None)
        output1 = ops.smooth_quant_matmul(quant_input, input_scale, quant_weight, weight_scale, q.dtype, bias)
        output_list = []
        block_size = seq // block_seq
        for i in range(block_seq):
            start = i * block_size
            end = seq if i == block_seq - 1 else (i + 1) * block_size
            attn_output = ops.flash_attention(q[:,start:end], k[:,:end], v[:,:end], None, None, None,
                None, None, end - start, end, softmax_scale, is_causal).flatten(-2, -1).flatten(0, 1)
            quant_input, input_scale = ops.per_token_smooth_quantize(attn_output, smooth, None)
            out = ops.smooth_quant_matmul(quant_input, input_scale, quant_weight, weight_scale, q.dtype, bias)
            output_list.append(out)
        output2 = torch.cat(output_list, dim=0)
        output2 = output2.reshape(block_seq, batch, block_size, hidden_size).transpose(0, 1).reshape(batch*seq, hidden_size)
        assertTensorsEqual(output1.cpu().float(), output2.cpu().float(), 0.0045, use_MSE=True, use_RAE=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_flash_attn_sq_mm_allreduce(self):
        world_size = min(torch_mlu.mlu.device_count(), 8)
        batch, seq, head_num_q, head_num_kv, head_size, is_causal = 1, 8192, 64, 8, 128, True
        dtype = torch.bfloat16 if torch_mlu.mlu.is_bf16_supported() else torch.half
        for i in range(1):
            # seq = random.randint(1, 32768)
            # is_causal = bool(random.randint(0, 1))
            print("=============test[{}]: seq = {}, causal = {}===============".format(i, seq, is_causal))
            hidden_size = head_num_q * head_size
            softmax_scale = 1 / math.sqrt(head_size)
            qkv = torch.randn(batch, seq, head_num_q + 2 * head_num_kv, head_size, dtype=dtype).mlu()
            smooth = torch.zeros(hidden_size, dtype=torch.float).mlu() + 1.0
            bias = torch.randn(hidden_size, dtype=dtype).mlu()
            weight = torch.randn(hidden_size, hidden_size, dtype=dtype).mlu()
            quant_weight, weight_scale = QuantByRow(weight / smooth, 8)
            q = qkv[:, :, :head_num_q, :]
            k = qkv[:, :, head_num_q:head_num_q+head_num_kv, :]
            v = qkv[:, :, head_num_q+head_num_kv:, :]
            output1, output2 = flash_attn_sq_mm(q, k, v, smooth, quant_weight, weight_scale,
            bias, softmax_scale, is_causal, world_size)
            args = world_size, q, k, v, smooth, quant_weight, weight_scale, bias, softmax_scale, is_causal, output1
            mp.spawn(tp_flash_attn_sq_mm, args, nprocs=world_size, join=True)

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

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