# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from typing import TYPE_CHECKING, Literal, Optional

import torch
from compressed_tensors.offload.cache.base import OffloadCache
from compressed_tensors.offload.utils import send_tensors


if TYPE_CHECKING:
    from torch._prims_common import DeviceLikeType


class DeviceCache(OffloadCache):
    """
    Handles offloading and onloading tensors from/to device memory. Onloading
    tensors is typically a no-op (except onload device has been modified).
    """

    def __init__(
        self,
        onload_device: "DeviceLikeType",
        offload_device: Optional["DeviceLikeType | Literal['disk']"] = None,
    ):
        super().__init__(onload_device, offload_device=offload_device)
        if offload_device is not None:
            self.offload_device = torch.device(offload_device)
        else:
            self.offload_device = self.onload_device

    def onload(self, offloaded: torch.Tensor | None) -> torch.Tensor | None:
        """
        Typically a no op, except when onload device has been modified

        :param key: device tensor to onload
        :return: device tensor
        """
        # move because onload_device might be modified after init
        return send_tensors(offloaded, device=self.onload_device, copy=False)

    def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None:
        """
        Offload a tensor to the device

        :param value: tensor on any device
        :return: tensor on device
        """
        return send_tensors(tensor, device=self.offload_device, copy=False)

    @torch.no_grad()
    def update_offload(self, offloaded: torch.Tensor, data: torch.Tensor | None):
        """
        Update the device value with new data

        :param offloaded: device tensor to update
        :param data: new data to copy from
        """
        if data is not None:
            offloaded.copy_(data)
