torch.accelerator 的源代码

r""
此包引入了对当前 python 中:ref:`加速器`的支持。
""

from 打字 导入 可选
from typing_extensions 导入 已弃用

导入 火炬

from _utils 导入 _device_t, _get_device_index


全部 = [
    current_accelerator,
    current_device_idx,  # 已弃用
    current_device_index,
    "当前流",
    设备数量,
    "是否可用",
    "设置设备索引",  # 已弃用
    "设置设备索引",
    "设置流",
    "同步",
]


[文档]def 设备数量() -> int:
返回当前可用的 :ref:`加速器` 的数量。

返回值：
int: 当前可用的 :ref:`加速器` 的数量。
如果没有可用的加速器，则返回 0。
"""
返回加速器的设备数量


[文档]def is_available() -> bool:
检查当前加速器在运行时是否可用：它已构建，所有
必需的驱动程序可用且至少有一个设备可见。
详细信息请参阅：:ref:`加速器`。

返回：
bool：一个布尔值，表示是否存在可用的：ref:`加速器`。

示例::

>>> 断言 torch.accelerator.is_available() "未检测到可用的加速器。"
"""
# 为什么不就像其他 is_available 调用一样检查 "device_count() > 0" 呢？
因为像 CUDA 这样的设备有 Python 实现的 is_available，它是非破坏性的，并且一些功能如 Dataloader 依赖于它。
因此我们在这里小心地将加速器的 Python 版本委托给这里。
所以我们在这里小心地委托给加速器的 Python 版本。
acc = current_accelerator()
如果 acc 为 None：
返回 False

mod = torch.get_device_module(acc)
返回 mod.is_available()


[文档]def current_accelerator(check_available: bool = False) -> Optional[torch.device]:
r"""返回编译时可用加速器的设备。
如果编译时没有可用加速器，则返回 None。
详细信息请参阅 :ref:`加速器`。

Args:
check_available (bool, optional): 如果为 True，将还会进行运行时检查，以查看
是否设备支持 :func:`torch.accelerator.is_available`，这将在编译时检查之上。
check.
默认：``False``

返回：
torch.device：返回当前加速器作为 :class:`torch.device`。

.. 注意：返回的 :class:`torch.device` 的索引将为 ``None``，请使用
`torch.accelerator.current_device_index` 获取当前使用的索引。

示例::

>>> # xdoctest:
>>> # 如果有可用的加速器，将模型发送到它
>>> model = torch.nn.Linear(2, 2)
>>> if (current_device := current_accelerator(check_available=True)) is not None:
>>>     model.to(current_device)
"""
if (acc := torch._C._accelerator_getAccelerator()) is not None:
if (not check_available) or (check_available and is_available()):
return acc
return None


[文档]def current_device_index() -> int:
r"""返回当前选定的设备的索引，该设备用于当前 :ref:`加速器`。

返回：
int: 当前选定设备的索引。
"``"
返回 torch._C._accelerator_getDeviceIndex()


current_device_idx = 已弃用(
    "请使用 `current_device_index`。",
    类别=FutureWarning,
)(current_device_index)


[文档]def 设置设备索引(device: _device_t, /) -> None:
设置当前设备索引为指定设备。

参数：
device (:class:`torch.device`, str, int): 必须匹配当前
ref:`加速器` 设备类型的给定设备。

.. 注意:: 如果此设备索引为负，则此函数为空操作。
"""
device_index = _get_device_index(device)
torch._C._accelerator_setDeviceIndex(device_index)


设置设备索引 = 已弃用(
    "请使用 `set_device_index` 代替。",
    类别=FutureWarning,
)(设置设备索引)


[文档]def current_stream(device: _device_t = None, /) -> torch.Stream:
返回给定设备的当前流。

Args:
设备 (:class:`torch.device`, str, int, 可选): 必须匹配当前环境的给定设备
ref:`加速器` 设备类型。如果未指定，
则默认使用 :func:`torch.accelerator.current_device_index`。

返回值：
torch.Stream：给定设备当前选定的流。
"""
device_index = _get_device_index(device, True)
return torch._C._accelerator_getStream(device_index)


[文档]def set_stream(stream: torch.Stream) -> None:
r"""设置当前流为指定的流。

Args:
stream (torch.Stream): 必须匹配当前 :ref:`加速器` 设备类型的给定流。

.. 注意:: 此函数将当前设备索引设置为给定流的设备索引。
"""
torch._C._accelerator_setStream(stream)


[文档]def 同步(device: _device_t = None, /) -> None:
等待给定设备上所有流中的所有内核完成。

参数：
device (:class:`torch.device`, str, int, optional): 要同步的设备。它必须匹配
当前 :ref:`加速器` 设备类型。如果没有提供，
默认情况下使用 :func:`torch.accelerator.current_device_index`。

.. note:: 如果当前 :ref:`accelerator` 未初始化，此函数为空操作。

示例::

>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
>>> assert torch.accelerator.is_available() "未检测到可用的加速器。"
>>> start_event = torch.Event(enable_timing=True)
>>> end_event = torch.Event(enable_timing=True)
>>> start_event.record()
>>> tensor = torch.randn(100, device=torch.accelerator.current_accelerator())
>>> sum = torch.sum(tensor)
>>> end_event.record()
>>> torch.accelerator.synchronize()
>>> elapsed_time_ms = start_event.elapsed_time(end_event)
"""
device_index = _get_device_index(device, True)
torch._C._accelerator_synchronizeDevice(device_index)