import torch
from typing import Union
from vllm.distributed.parallel_state import GroupCoordinator, supports_custom_op
from vllm_mlu.mlu_hijack_utils import MluHijackObject

def vllm__distributed__parallel_state__GroupCoordinator__all_reduce(
        self, input_: torch.Tensor,
        async_op: bool = False) -> Union[torch.distributed.Work, torch.Tensor]:
    """
    User-facing all-reduce function before we actually call the
    all-reduce operation.

    We need this because Dynamo does not support passing an arbitrary
    object (`self` in this case) to a custom op. We need to pass the
        group name as a string, and then look up the group coordinator from
        the group name, dispatch the all-reduce operation to the group
        coordinator.

    In addition, PyTorch custom ops do not support mutation or returning
    a new tensor in the same op. So we need to figure out if the op is
    in-place or out-of-place ahead of time.
    """
    # Bypass the function if we are using only 1 GPU.
    if self.world_size == 1:
        return input_

    if input_.is_cpu:
        import intel_extension_for_pytorch as ipex
        ipex.distributed.all_reduce(input_, group=self.device_group)
        return input_

    if not supports_custom_op():
        self._all_reduce_in_place(input_)
        return input_

    if self.tpu_communicator is not None and \
        not self.tpu_communicator.disabled:
        # TPU handles Dynamo with its own logic.
        return self.tpu_communicator.all_reduce(input_)

    if self.hpu_communicator is not None and \
        not self.hpu_communicator.disabled:
        return self.hpu_communicator.all_reduce(input_)

    if self.xpu_communicator is not None and \
            not self.xpu_communicator.disabled:
        return self.xpu_communicator.all_reduce(input_)

    if self.ca_comm is not None and \
        not self.ca_comm.disabled and \
            self.ca_comm.should_custom_ar(input_):
        return torch.ops.vllm.outplace_all_reduce(
            input_, group_name=self.unique_name)
    else:
        '''
        =============================
        Modify by custom vllm_mlu
        =============================
        @brief: use async all reduce when preload weights.
        '''
        handle = torch.distributed.all_reduce(input_, group=self.device_group,
                                              async_op=async_op)
        if async_op:
            return handle
        '''
        ==================
        End of custom MLU Hijack
        ==================
        '''
        return input_


MluHijackObject.apply_hijack(GroupCoordinator,
                             GroupCoordinator.all_reduce,
                             vllm__distributed__parallel_state__GroupCoordinator__all_reduce)