torch.cpu 的源代码

r"""
此包实现了在 ``torch.cuda`` 中找到的抽象，以便于编写与设备无关的代码。
"""

from contextlib import AbstractContextManager
from typing import Any, Optional, Union

import torch

from .. import device as _device
from . import amp

__all__ = [
    "is_available",
    "synchronize",
    "current_device",
    "current_stream",
    "stream",
    "set_device",
    "device_count",
    "Stream",
    "StreamContext",
    "Event",
]

_device_t = Union[_device, str, int, None]


def _is_cpu_support_vnni() -> bool:
    r"""返回一个布尔值，指示CPU是否支持VNNI。"""
    return torch._C._cpu._is_cpu_support_vnni()


def is_available() -> bool:
    r"""返回一个布尔值，指示CPU当前是否可用。

    注意：此函数仅存在以促进与设备无关的代码

    """
    return True


def synchronize(device: _device_t = None) -> None:
    r"""等待CPU设备上所有流中的所有内核完成。

    参数:
        device (torch.device 或 int, 可选): 忽略，只有一个CPU设备。

    注意：此函数仅存在以促进与设备无关的代码。
    """
    pass


class Stream:
    """
    注意：此类仅存在以促进与设备无关的代码
    """

    def __init__(self, priority: int = -1):
        pass

    def wait_stream(self, stream) -> None:
        pass


class Event:
    def query(self) -> bool:
        return True

    def record(self, stream=None):
        pass

    def synchronize(self):
        pass

    def wait(self, stream=None):
        pass


_default_cpu_stream = Stream()
_current_stream = _default_cpu_stream


def current_stream(device: _device_t = None) -> Stream:
    r"""返回给定设备的当前选定的 :class:`Stream`。

    参数:
        device (torch.device 或 int, 可选): 忽略。

    注意：此函数仅存在以促进与设备无关的代码

    """
    return _current_stream


class StreamContext(AbstractContextManager):
    r"""选择给定流的上下文管理器。

    注意：此类仅存在以促进与设备无关的代码

    """
    cur_stream: Optional[Stream]

    def __init__(self, stream):
        self.stream = stream
        self.prev_stream = _default_cpu_stream

    def __enter__(self):
        cur_stream = self.stream
        if cur_stream is None:
            return

        global _current_stream
        self.prev_stream = _current_stream
        _current_stream = cur_stream

    def __exit__(self, type: Any, value: Any, traceback: Any):
        cur_stream = self.stream
        if cur_stream is None:
            return

        global _current_stream
        _current_stream = self.prev_stream


def stream(stream: Stream) -> AbstractContextManager:
    r"""围绕上下文管理器 StreamContext 的包装器，
    选择给定的流。

    注意：此函数仅存在以促进与设备无关的代码
    """
    return StreamContext(stream)


def device_count() -> int:
    r"""返回CPU设备的数量（不是核心）。总是1。

    注意：此函数仅存在以促进与设备无关的代码
    """
    return 1


def set_device(device: <span class="n