from typing import Optional
import torch
from vllm.distributed.parallel_state import get_tp_group, get_tensor_model_parallel_group
from vllm.distributed import get_tensor_model_parallel_rank, split_tensor_along_last_dim
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor.layers.linear import UnquantizedLinearMethod, RowParallelLinear
from vllm import _mlu_ops as mlu_ops
from vllm_mlu._mlu_utils import *


def vllm__model_executor__layers__linear__UnquantizedLinearMethod__apply(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
    residual: Optional[torch.Tensor] = None
) -> torch.Tensor:
    beta = 1.0 if residual is not None else 0.0
    res_shape = x.shape[0:-1] + (layer.weight.shape[0], )
    '''
    =====================================================
    Modify by custom vllm_mlu
    =====================================================
    @brief: call parallel op if parallel_num is set
    '''
    if hasattr(self, 'parallel_num') and get_is_prompt():
        rank = get_tensor_model_parallel_rank()
        pg = get_tensor_model_parallel_group().device_group
        cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
        return mlu_ops.matmul_allreduce(cncl_comm, x.view(-1, x.shape[-1]), layer.weight, 
                                        bias, residual, 1.0, beta, self.parallel_num).view(res_shape)
    return mlu_ops.matmul(x.view(-1, x.shape[-1]), layer.weight, bias, residual, 'none', 1.0, beta).view(res_shape)
    '''
    =====================================================
    End of custom MLU Hijack
    =====================================================
    '''


def vllm__model_executor__layers__linear__RowParallelLinear__forward(
    self, 
    input_, 
    residual: Optional[torch.Tensor] = None
):
    if self.input_is_parallel:
        input_parallel = input_
    else:
        tp_rank = get_tensor_model_parallel_rank()
        splitted_input = split_tensor_along_last_dim(
            input_, num_partitions=self.tp_size)
        input_parallel = splitted_input[tp_rank].contiguous()

    # Matrix multiply.
    assert self.quant_method is not None
    # Only fuse bias add into GEMM for rank 0 (this ensures that
    # bias will not get added more than once in TP>1 case)
    bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
    residual_ = None if self.tp_rank > 0 else residual
    '''
    =====================================================
    Modify by custom vllm_mlu
    =====================================================
    @brief: abandon original reduce if parallel_num is set
    '''
    is_parallel_enable = hasattr(self.quant_method, 'parallel_num') and get_is_prompt()
    '''
    =====================================================
    End of custom MLU Hijack
    =====================================================
    '''
    output_parallel = self.quant_method.apply(self,
                                              input_parallel,
                                              bias=bias_,
                                              residual=residual_)
    '''
    =============================
    Modify by custom vllm_mlu
    =============================
    @brief: when preload_size is set, call GroupCoordinator.all_reduce() directly and
    use async_op to set all_reduce paralleled with preload
    '''
    if self.reduce_results and self.tp_size > 1 and not is_parallel_enable:
        if hasattr(self, 'preload_size') and self.preload_size > 0 and not self.is_prompt:
            handle = get_tp_group().all_reduce(output_parallel, async_op=True)
            _MB = 1 << 20
            mlu_ops.preload(self.preloaded_weights[0].data, self.preload_size * _MB)
            preloaded_weights_size = self.preloaded_weights[0].numel() * self.preloaded_weights[0].element_size()
            if preloaded_weights_size < (self.preload_size * _MB) and len(self.preloaded_weights) > 1:
                mlu_ops.preload(self.preloaded_weights[1].data, (self.preload_size * _MB) - preloaded_weights_size)
            handle.wait()
            output = output_parallel
        else:
            output = tensor_model_parallel_all_reduce(output_parallel)
    else:
        output = output_parallel
    '''
    =========================
    End of custom MLU Hijack
    =========================
    '''
    output_bias = self.bias if self.skip_bias_add else None

    return output, output_bias


MluHijackObject.undo_hijack(UnquantizedLinearMethod, 
                            UnquantizedLinearMethod.apply)
MluHijackObject.apply_hijack(UnquantizedLinearMethod,
                             UnquantizedLinearMethod.apply,
                             vllm__model_executor__layers__linear__UnquantizedLinearMethod__apply)
MluHijackObject.undo_hijack(RowParallelLinear, 
                            RowParallelLinear.forward)
MluHijackObject.apply_hijack(RowParallelLinear,
                             RowParallelLinear.forward,
                             vllm__model_executor__layers__linear__RowParallelLinear__forward)
