infer/modules/ipex/hijacks.py

import contextlib
import importlib
import torch
import intel_extension_for_pytorch as ipex  # pylint: disable=import-error, unused-import

# pylint: disable=protected-access, missing-function-docstring, line-too-long, unnecessary-lambda, no-else-return


class CondFunc:  # pylint: disable=missing-class-docstring
    def __new__(cls, orig_func, sub_func, cond_func):
        self = super(CondFunc, cls).__new__(cls)
        if isinstance(orig_func, str):
            func_path = orig_func.split(".")
            for i in range(len(func_path) - 1, -1, -1):
                try:
                    resolved_obj = importlib.import_module(".".join(func_path[:i]))
                    break
                except ImportError:
                    pass
            for attr_name in func_path[i:-1]:
                resolved_obj = getattr(resolved_obj, attr_name)
            orig_func = getattr(resolved_obj, func_path[-1])
            setattr(
                resolved_obj,
                func_path[-1],
                lambda *args, **kwargs: self(*args, **kwargs),
            )
        self.__init__(orig_func, sub_func, cond_func)
        return lambda *args, **kwargs: self(*args, **kwargs)

    def __init__(self, orig_func, sub_func, cond_func):
        self.__orig_func = orig_func
        self.__sub_func = sub_func
        self.__cond_func = cond_func

    def __call__(self, *args, **kwargs):
        if not self.__cond_func or self.__cond_func(self.__orig_func, *args, **kwargs):
            return self.__sub_func(self.__orig_func, *args, **kwargs)
        else:
            return self.__orig_func(*args, **kwargs)


_utils = torch.utils.data._utils


def _shutdown_workers(self):
    if (
        torch.utils.data._utils is None
        or torch.utils.data._utils.python_exit_status is True
        or torch.utils.data._utils.python_exit_status is None
    ):
        return
    if hasattr(self, "_shutdown") and not self._shutdown:
        self._shutdown = True
        try:
            if hasattr(self, "_pin_memory_thread"):
                self._pin_memory_thread_done_event.set()
                self._worker_result_queue.put((None, None))
                self._pin_memory_thread.join()
                self._worker_result_queue.cancel_join_thread()
                self._worker_result_queue.close()
            self._workers_done_event.set()
            for worker_id in range(len(self._workers)):
                if self._persistent_workers or self._workers_status[worker_id]:
                    self._mark_worker_as_unavailable(worker_id, shutdown=True)
            for w in self._workers:  # pylint: disable=invalid-name
                w.join(timeout=torch.utils.data._utils.MP_STATUS_CHECK_INTERVAL)
            for q in self._index_queues:  # pylint: disable=invalid-name
                q.cancel_join_thread()
                q.close()
        finally:
            if self._worker_pids_set:
                torch.utils.data._utils.signal_handling._remove_worker_pids(id(self))
                self._worker_pids_set = False
            for w in self._workers:  # pylint: disable=invalid-name
                if w.is_alive():
                    w.terminate()


class DummyDataParallel(
    torch.nn.Module
):  # pylint: disable=missing-class-docstring, unused-argument, too-few-public-methods
    def __new__(
        cls, module, device_ids=None, output_device=None, dim=0
    ):  # pylint: disable=unused-argument
        if isinstance(device_ids, list) and len(device_ids) > 1:
            print("IPEX backend doesn't support DataParallel on multiple XPU devices")
        return module.to("xpu")


def return_null_context(*args, **kwargs):  # pylint: disable=unused-argument
    return contextlib.nullcontext()


def check_device(device):
    return bool(
        (isinstance(device, torch.device) and device.type == "cuda")
        or (isinstance(device, str) and "cuda" in device)
        or isinstance(device, int)
    )


def return_xpu(device):
    return (
        f"xpu:{device[-1]}"
        if isinstance(device, str) and ":" in device
        else (
            f"xpu:{device}"
            if isinstance(device, int)
            else torch.device("xpu") if isinstance(device, torch.device) else "xpu"
        )
    )


def ipex_no_cuda(orig_func, *args, **kwargs):
    torch.cuda.is_available = lambda: False
    orig_func(*args, **kwargs)
    torch.cuda.is_available = torch.xpu.is_available


original_autocast = torch.autocast


def ipex_autocast(*args, **kwargs):
    if len(args) > 0 and args[0] == "cuda":
        return original_autocast("xpu", *args[1:], **kwargs)
    else:
        return original_autocast(*args, **kwargs)


original_torch_cat = torch.cat


def torch_cat(tensor, *args, **kwargs):
    if len(tensor) == 3 and (
        tensor[0].dtype != tensor[1].dtype or tensor[2].dtype != tensor[1].dtype
    ):
        return original_torch_cat(
            [tensor[0].to(tensor[1].dtype), tensor[1], tensor[2].to(tensor[1].dtype)],
            *args,
            **kwargs,
        )
    else:
        return original_torch_cat(tensor, *args, **kwargs)


original_interpolate = torch.nn.functional.interpolate


def interpolate(
    tensor,
    size=None,
    scale_factor=None,
    mode="nearest",
    align_corners=None,
    recompute_scale_factor=None,
    antialias=False,
):  # pylint: disable=too-many-arguments
    if antialias or align_corners is not None:
        return_device = tensor.device
        return_dtype = tensor.dtype
        return original_interpolate(
            tensor.to("cpu", dtype=torch.float32),
            size=size,
            scale_factor=scale_factor,
            mode=mode,
            align_corners=align_corners,
            recompute_scale_factor=recompute_scale_factor,
            antialias=antialias,
        ).to(return_device, dtype=return_dtype)
    else:
        return original_interpolate(
            tensor,
            size=size,
            scale_factor=scale_factor,
            mode=mode,
            align_corners=align_corners,
            recompute_scale_factor=recompute_scale_factor,
            antialias=antialias,
        )


original_linalg_solve = torch.linalg.solve


def linalg_solve(A, B, *args, **kwargs):  # pylint: disable=invalid-name
    if A.device != torch.device("cpu") or B.device != torch.device("cpu"):
        return_device = A.device
        return original_linalg_solve(A.to("cpu"), B.to("cpu"), *args, **kwargs).to(
            return_device
        )
    else:
        return original_linalg_solve(A, B, *args, **kwargs)


def ipex_hijacks():
    CondFunc(
        "torch.Tensor.to",
        lambda orig_func, self, device=None, *args, **kwargs: orig_func(
            self, return_xpu(device), *args, **kwargs
        ),
        lambda orig_func, self, device=None, *args, **kwargs: check_device(device),
    )
    CondFunc(
        "torch.Tensor.cuda",
        lambda orig_func, self, device=None, *args, **kwargs: orig_func(
            self, return_xpu(device), *args, **kwargs
        ),
        lambda orig_func, self, device=None, *args, **kwargs: check_device(device),
    )
    CondFunc(
        "torch.empty",
        lambda orig_func, *args, device=None, **kwargs: orig_func(
            *args, device=return_xpu(device), **kwargs
        ),
        lambda orig_func, *args, device=None, **kwargs: check_device(device),
    )
    CondFunc(
        "torch.load",
        lambda orig_func, *args, map_location=None, **kwargs: orig_func(
            *args, return_xpu(map_location), **kwargs
        ),
        lambda orig_func, *args, map_location=None, **kwargs: map_location is None
        or check_device(map_location),
    )
    CondFunc(
        "torch.randn",
        lambda orig_func, *args, device=None, **kwargs: orig_func(
            *args, device=return_xpu(device), **kwargs
        ),
        lambda orig_func, *args, device=None, **kwargs: check_device(device),
    )
    CondFunc(
        "torch.ones",
        lambda orig_func, *args, device=None, **kwargs: orig_func(
            *args, device=return_xpu(device), **kwargs
        ),
        lambda orig_func, *args, device=None, **kwargs: check_device(device),
    )
    CondFunc(
        "torch.zeros",
        lambda orig_func, *args, device=None, **kwargs: orig_func(
            *args, device=return_xpu(device), **kwargs
        ),
        lambda orig_func, *args, device=None, **kwargs: check_device(device),
    )
    CondFunc(
        "torch.tensor",
        lambda orig_func, *args, device=None, **kwargs: orig_func(
            *args, device=return_xpu(device), **kwargs
        ),
        lambda orig_func, *args, device=None, **kwargs: check_device(device),
    )
    CondFunc(
        "torch.linspace",
        lambda orig_func, *args, device=None, **kwargs: orig_func(
            *args, device=return_xpu(device), **kwargs
        ),
        lambda orig_func, *args, device=None, **kwargs: check_device(device),
    )

    CondFunc(
        "torch.Generator",
        lambda orig_func, device=None: torch.xpu.Generator(device),
        lambda orig_func, device=None: device is not None
        and device != torch.device("cpu")
        and device != "cpu",
    )

    CondFunc(
        "torch.batch_norm",
        lambda orig_func, input, weight, bias, *args, **kwargs: orig_func(
            input,
            (
                weight
                if weight is not None
                else torch.ones(input.size()[1], device=input.device)
            ),
            (
                bias
                if bias is not None
                else torch.zeros(input.size()[1], device=input.device)
            ),
            *args,
            **kwargs,
        ),
        lambda orig_func, input, *args, **kwargs: input.device != torch.device("cpu"),
    )
    CondFunc(
        "torch.instance_norm",
        lambda orig_func, input, weight, bias, *args, **kwargs: orig_func(
            input,
            (
                weight
                if weight is not None
                else torch.ones(input.size()[1], device=input.device)
            ),
            (
                bias
                if bias is not None
                else torch.zeros(input.size()[1], device=input.device)
            ),
            *args,
            **kwargs,
        ),
        lambda orig_func, input, *args, **kwargs: input.device != torch.device("cpu"),
    )

    # Functions with dtype errors:
    CondFunc(
        "torch.nn.modules.GroupNorm.forward",
        lambda orig_func, self, input: orig_func(
            self, input.to(self.weight.data.dtype)
        ),
        lambda orig_func, self, input: input.dtype != self.weight.data.dtype,
    )
    CondFunc(
        "torch.nn.modules.linear.Linear.forward",
        lambda orig_func, self, input: orig_func(
            self, input.to(self.weight.data.dtype)
        ),
        lambda orig_func, self, input: input.dtype != self.weight.data.dtype,
    )
    CondFunc(
        "torch.nn.modules.conv.Conv2d.forward",
        lambda orig_func, self, input: orig_func(
            self, input.to(self.weight.data.dtype)
        ),
        lambda orig_func, self, input: input.dtype != self.weight.data.dtype,
    )
    CondFunc(
        "torch.nn.functional.layer_norm",
        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs: orig_func(
            input.to(weight.data.dtype), normalized_shape, weight, *args, **kwargs
        ),
        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs: weight
        is not None
        and input.dtype != weight.data.dtype,
    )

    # Diffusers Float64 (ARC GPUs doesn't support double or Float64):
    if not torch.xpu.has_fp64_dtype():
        CondFunc(
            "torch.from_numpy",
            lambda orig_func, ndarray: orig_func(ndarray.astype("float32")),
            lambda orig_func, ndarray: ndarray.dtype == float,
        )

    # Broken functions when torch.cuda.is_available is True:
    CondFunc(
        "torch.utils.data.dataloader._BaseDataLoaderIter.__init__",
        lambda orig_func, *args, **kwargs: ipex_no_cuda(orig_func, *args, **kwargs),
        lambda orig_func, *args, **kwargs: True,
    )

    # Functions that make compile mad with CondFunc:
    torch.utils.data.dataloader._MultiProcessingDataLoaderIter._shutdown_workers = (
        _shutdown_workers
    )
    torch.nn.DataParallel = DummyDataParallel
    torch.autocast = ipex_autocast
    torch.cat = torch_cat
    torch.linalg.solve = linalg_solve
    torch.nn.functional.interpolate = interpolate
    torch.backends.cuda.sdp_kernel = return_null_context