torch.backends.mkldnn 的源代码

import sys
from contextlib import contextmanager

from typing import TYPE_CHECKING

import torch
from torch.backends import __allow_nonbracketed_mutation, ContextProp, PropModule


def is_available():
    r"""返回 PyTorch 是否使用 MKL-DNN 支持构建。"""
    return torch._C._has_mkldnn


VERBOSE_OFF = 0
VERBOSE_ON = 1
VERBOSE_ON_CREATION = 2


class verbose:
    """
    按需的 oneDNN（前身为 MKL-DNN）详细功能。

    为了更容易调试性能问题，oneDNN 可以转储详细消息，其中包含有关内核大小、输入数据大小和执行内核时的执行持续时间等信息。详细功能可以通过名为 `DNNL_VERBOSE` 的环境变量调用。然而，这种方法会在所有步骤中转储消息。这些是大量的详细消息。此外，为了调查性能问题，通常只需要一个迭代的详细消息就足够了。此按需详细功能使得可以控制详细消息转储的范围。在以下示例中，详细消息将仅在第二次推理时转储。

    .. highlight:: python
    .. code-block:: python

        import torch
        model(data)
        with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON):
            model(data)

    参数:
        level: 详细级别
            - ``VERBOSE_OFF``: 禁用详细
            - ``VERBOSE_ON``:  启用详细
            - ``VERBOSE_ON_CREATION``: 启用详细，包括 oneDNN 内核创建
    """

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

    def __enter__(self):
        if self.level == VERBOSE_OFF:
            return
        st = torch._C._verbose.mkldnn_set_verbose(self.level)
        assert (
            st
        ), "Failed to set MKLDNN into verbose mode. Please consider to disable this verbose scope."
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        torch._C._verbose.mkldnn_set_verbose(VERBOSE_OFF)
        return False


def set_flags(_enabled):
    orig_flags = (torch._C._get_mkldnn_enabled(),)
    torch._C._set_mkldnn_enabled(_enabled)
    return orig_flags


@contextmanager
def flags(enabled=False):
    with __allow_nonbracketed_mutation():
        orig_flags = set_flags(enabled)
    try:
        yield
    finally:
        with __allow_nonbracketed_mutation():
            set_flags(orig_flags[0])


class MkldnnModule(PropModule):
    def __init__(self, m, name):
        super().__init__(m, name)

    enabled = ContextProp(torch._C._get_mkldnn_enabled, torch._C._set_mkldnn_enabled)


if TYPE_CHECKING:
    enabled: ContextProp


# Cool stuff from torch/backends/cudnn/__init__.py and
# https://stackoverflow.com/questions/2447353/getattr-on-a-module/7668273#7668273
sys.modules[__name__] = MkldnnModule(sys.modules[__name__], __name__)