torch.compiler 的源代码

import torch
from typing import List

__all__ = [
    "compile",
    "assume_constant_result",
    "reset",
    "allow_in_graph",
    "list_backends",
    "disable",
    "cudagraph_mark_step_begin",
    "wrap_numpy",
    "is_compiling",
    "is_dynamo_compiling",
]

def compile(*args, **kwargs):
    """
    有关此函数的参数的详细信息，请参阅 :func:`torch.compile`。
    """
    return torch.compile(*args, **kwargs)

def reset() -> None:
    """
    此函数清除所有编译缓存并将系统恢复到其初始状态。
    建议调用此函数，特别是在使用 `torch.compile(...)` 等操作后，以确保在另一个不相关的编译之前处于干净状态。
    """
    import torch._dynamo

    torch._dynamo.reset()

def allow_in_graph(fn):
    """
    自定义编译将包含在生成的图中的函数。
    它绕过了对符号化Python代码的所有自省，直接将其写入图中。
    如果 fn 是一个可调用对象的列表或元组，它将递归地应用 :func:`allow_in_graph()`
    到每个函数，并返回包含修改后的函数的新列表或元组。

    参数:
        fn: 表示要包含在图中的函数的可调用对象。

    .. 警告::

        :func:`allow_in_graph` 完全绕过了装饰函数的 TorchDynamo
        跳过了所有 TorchDynamo 安全检查（图中断、处理闭包等）。
        因此，使用 :func:`allow_in_graph` 时要非常小心，因为子系统
        如 AOT Autograd 依赖于 torchdynamo。
        如果不小心，这可能会导致正确性和非常难以调试的问题。

    """
    import torch._dynamo

    return torch._dynamo.allow_in_graph(fn)


def list_backends(exclude_tags=("debug", "experimental")) -> List[str]:
    """
    返回可以传递给 `torch.compile(..., backend="name")` 的有效字符串。

    参数:
        exclude_tags(可选): 表示要排除的标签的字符串元组。
    """
    import torch._dynamo

    return torch._dynamo.list_backends(exclude_tags)

def assume_constant_result(fn):
    """
    此函数用于将函数 `fn` 标记为具有常量结果。
    这允许编译器优化掉您的函数。
    返回相同的函数 `fn`。

    参数:
        fn: 要标记为具有常量结果的函数。

    .. 警告::
        `assume_constant_result` 如果无效可能会导致安全性和正确性问题，:func:`torch.compile`
        不会尝试验证常量假设是否正确。

    """
    import torch._dynamo

    return torch._dynamo.assume_constant_result(fn)

def disable(fn=None, recursive=True):
    """
    此函数提供了一个装饰器和一个上下文管理器来禁用对函数的编译。
    它还提供了递归禁用调用函数的功能。

    参数:
        fn (可选): 要禁用的函数。
        recursive (可选): 一个布尔值，指示禁用是否应递归。
    """
    import torch._dynamo

    return torch._dynamo.disable(fn, recursive)

def cudagraph_mark_step_begin():
    """
    指示推理或训练的新迭代即将开始。

    CUDA 图将释放先前迭代的张量。每次调用 torch.compile 时都会启动一个新的迭代，
    只要没有未调用的待处理反向传播。

    如果该启发式方法是错误的，例如在以下示例中，请使用此 API 手动标记。

    .. 代码块:: python

        @torch.compile(mode="reduce-overhead")
        def rand_foo():
            return torch.rand([4], device="cuda")

        for _ in range(5):
            torch.compiler.cudagraph_mark_step_begin()
            rand_foo() + rand_foo()

    有关更多详细信息，请参阅 `torch.compiler_cudagraph_trees `__
    """
    from torch._inductor import cudagraph_trees

    cudagraph_trees.mark_step_begin()

def wrap_numpy(fn):
    r"""装饰器，将一个从 ``np.ndarray`` 到 ``np.ndarray`` 的函数转换为一个从 ``torch.Tensor`` 到 ``torch.Tensor`` 的函数。
    它设计用于 :func:`torch.compile` 与 ``fullgraph=True`` 一起使用。它允许将 NumPy 函数编译为 PyTorch 函数。
    这使您可以在 CUDA 上运行 NumPy 代码或计算其梯度。

    .. 注意::

        此装饰器在没有 :func:`torch.compile` 的情况下不起作用。

    示例::

        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
        >>> # 将 NumPy 函数编译为 Tensor -> Tensor 函数
        >>> @torch.compile(fullgraph=True)
        >>> @torch.compiler.wrap_numpy
        >>> def fn(a: np.ndarray):
        >>>     return np.sum(a * a)
        >>> # 使用 CUDA 上的 Tensors 执行 NumPy 函数并计算梯度
        >>> x = torch.arange(6, dtype=torch.float32, device="cuda", requires_grad=True)
        >>> out = fn(x)
        >>> out.backward()
        >>> print(x.grad)
        tensor([ 0.,  2.,  4.,  6.,  8., 10.], device='cuda:0')
    """
    from torch._dynamo.external_utils import wrap_numpy as wrap
    return wrap(fn)

_is_compiling_flag: bool = False

<a class="viewcode-back" href="../../generated