torch.profiler.itt 的源代码

from contextlib import contextmanager

try:
    from torch._C import _itt
except ImportError:

    class _ITTStub:
        @staticmethod
        def _fail(*args, **kwargs):
            raise RuntimeError(
                "未安装ITT功能。您确定安装了ITT构建吗？"
            )

        @staticmethod
        def is_available():
            return False

        rangePush = _fail
        rangePop = _fail
        mark = _fail

    _itt = _ITTStub()  # type: ignore[assignment]


__all__ = ["is_available", "range_push", "range_pop", "mark", "range"]


def is_available():
    """
    检查ITT功能是否可用
    """
    return _itt.is_available()


def range_push(msg):
    """
    将一个范围推入嵌套范围跨度的堆栈。返回开始范围的基于零的深度。

    参数:
        msg (str): 与范围关联的ASCII消息
    """
    return _itt.rangePush(msg)


def range_pop():
    """
    从嵌套范围跨度的堆栈中弹出一个范围。返回结束范围的基于零的深度。
    """
    return _itt.rangePop()


def mark(msg):
    """
    描述在某个时间点发生的瞬时事件。

    参数:
        msg (str): 与事件关联的ASCII消息。
    """
    return _itt.mark(msg)


@contextmanager
def range(msg, *args, **kwargs):
    """
    上下文管理器/装饰器，在作用域开始时推送一个ITT范围，在结束时弹出它。如果给出额外参数，它们将作为参数传递给msg.format()。

    参数:
        msg (str): 与范围关联的消息
    """
    range_push(msg.format(*args, **kwargs))
    try:
        yield
    finally:
        range_pop()