torch.backends.cuda 的源代码

import contextlib
import warnings

from typing import Union

import torch

__all__ = [
    "is_built",
    "cuFFTPlanCacheAttrContextProp",
    "cuFFTPlanCache",
    "cuFFTPlanCacheManager",
    "cuBLASModule",
    "preferred_linalg_library",
    "cufft_plan_cache",
    "matmul",
    "SDPBackend",
    "SDPAParams",
    "enable_cudnn_sdp",
    "cudnn_sdp_enabled",
    "enable_flash_sdp",
    "flash_sdp_enabled",
    "enable_mem_efficient_sdp",
    "mem_efficient_sdp_enabled",
    "math_sdp_enabled",
    "enable_math_sdp",
    "can_use_flash_attention",
    "can_use_efficient_attention",
    "sdp_kernel",
]


def is_built():
    r"""
    返回 PyTorch 是否构建了 CUDA 支持。

    请注意，这并不一定意味着 CUDA 是可用的；只是如果这个 PyTorch 二进制文件在具有工作 CUDA 驱动程序和设备的机器上运行，我们将能够使用它。
    """
    return torch._C._has_cuda


class cuFFTPlanCacheAttrContextProp:
    # 类似于常规的 ContextProp，但使用调用对象的 `.device_index` 属性作为 getter 和 setter 的第一个参数。
    def __init__(self, getter, setter):
        self.getter = getter
        self.setter = setter

    def __get__(self, obj, objtype):
        return self.getter(obj.device_index)

    def __set__(self, obj, val):
        if isinstance(self.setter, str):
            raise RuntimeError(self.setter)
        self.setter(obj.device_index, val)


class cuFFTPlanCache:
    r"""
    表示特定 `device_index` 的特定计划缓存。

    属性 `size` 和 `max_size`，以及方法 `clear`，可以获取和/或更改 C++ cuFFT 计划缓存的属性。
    """

    def __init__(self, device_index):
        self.device_index = device_index

    size = cuFFTPlanCacheAttrContextProp(
        torch._cufft_get_plan_cache_size,
        ".size 是一个只读属性，显示当前缓存中的计划数量。要更改缓存容量，请设置 cufft_plan_cache.max_size。",
    )

    max_size = cuFFTPlanCacheAttrContextProp(
        torch._cufft_get_plan_cache_max_size, torch._cufft_set_plan_cache_max_size
    )

    def clear(self):
        return torch._cufft_clear_plan_cache(self.device_index)


class cuFFTPlanCacheManager:
    r"""
    表示所有 cuFFT 计划缓存，当索引时返回给定设备的 cuFFTPlanCache。

    最后，当直接使用此对象作为 `cuFFTPlanCache` 对象（例如，设置 `.max_size` 属性）时，将使用当前设备的 cuFFT 计划缓存。
    """

    __initialized = False

    def __init__(self):
        self.caches = []
        self.__initialized = True

    def __getitem__(self, device):
        index = torch.cuda._utils._get_device_index(device)
        if index < 0 or index >= torch.cuda.device_count():
            raise RuntimeError(
                f"cufft_plan_cache: 预期 0 <= 设备索引 < {torch.cuda.device_count()}，但得到设备索引 {index}"
            )
        if len(self.caches) == 0:
            self.caches.extend(
                cuFFTPlanCache(index) for index in range(torch.cuda.device_count())
            )
        return self.caches[index]

    def __getattr__(self, name):
        return getattr(self[torch.cuda.current_device()], name)

    def __setattr__(self, name, value):
        if self.__initialized:
            return setattr(self[torch.cuda.current_device()], name, value)
        else:
            return super().__setattr__(name<span class="p