# Owner(s): ["module: mkldnn"]
import sys
import torch
import unittest
import itertools
import torch.nn as nn
from functools import wraps
from concurrent import futures
import torch.nn.functional as F
import torch.fx.experimental.optimization as optimization
from torch.testing._internal.jit_utils import JitTestCase
from torch.testing._internal.common_utils import run_tests, TEST_SCIPY, IS_WINDOWS, IS_MACOS
from torch.testing._internal.common_device_type import (
    instantiate_device_type_tests,
    onlyCPU,
    dtypes
)

# We use this wrapper to run UTs of TorchVision models because of a memory-leak
# issue with JIT tracing that causes traced model objects to persist in the
# memory. Ref: https://github.com/pytorch/pytorch/issues/35600
# Memory requirement for running these UTs was thus increasing cumulatively, and
# invoked the Linux kernel OOM killer on linux.2xlarge PyTorch CI runners, which
# only have 16 GB RAM. Cumulatively, these UTs had been using more than 14 GB
# memory (as per psutils). So now we run each TorchVision model UTs in separate processes.
def separate_process(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        with futures.ProcessPoolExecutor() as executor:
            future = executor.submit(func, *args, **kwargs)
            futures.wait([future])
    return wrapper

def is_avx512_supported():
    if sys.platform != 'linux':
        return False
    with open("/proc/cpuinfo", encoding="ascii") as f:
        lines = f.read()
    return "avx512" in lines

IS_AVX512_UNSUPPORTED = not is_avx512_supported()

LLGA_FUSION_GROUP = 'prim::oneDNNFusionGroup'
LLGA_NOT_ENABLED = not torch.backends.mkldnn.is_available() or IS_WINDOWS or IS_MACOS

def warmup_forward(f, *args, profiling_count=3):
    for i in range(profiling_count):
        results = f(*args)

    return results

class JitLlgaTestCase(JitTestCase):

    def setUp(self):
        # PyTorch has divergent op support for AMP in JIT & eager modes
        # so we disable AMP for JIT & leverage eager-mode AMP.
        # Ref: https://github.com/pytorch/pytorch/issues/75956
        self.original_autocast_mode = torch._C._jit_set_autocast_mode(False)
        torch.jit.enable_onednn_fusion(True)

    def tearDown(self):
        torch.jit.enable_onednn_fusion(False)
        torch._C._jit_set_autocast_mode(self.original_autocast_mode)

    def checkTrace(self, m, x, dtype=torch.float32, *args, **kwargs):
        if isinstance(m, torch.nn.Module):
            m.eval()
        with torch.no_grad(), torch._jit_internal._disable_emit_hooks():
            if dtype == torch.bfloat16:
                # We rely upon eager-mode AMP support for BF16
                with torch.cpu.amp.autocast(cache_enabled=False, dtype=torch.bfloat16):
                    traced = torch.jit.trace(m, x)
                    if isinstance(m, torch.nn.Module):
                        traced = torch.jit.freeze(traced)
                    warmup_forward(traced, *x)
                    ref_o = m(*x)
                    fwd_graph = traced.graph_for(*x)
            else:
                traced = torch.jit.trace(m, x)
                if isinstance(m, torch.nn.Module):
                    traced = torch.jit.freeze(traced)
                warmup_forward(traced, *x)
                ref_o = m(*x)
                fwd_graph = traced.graph_for(*x)

            jit_o = traced(*x)
            self.assertEqual(jit_o, ref_o)
            return traced, fwd_graph


    def assertFused(self, graph, fused_patterns):
        for pat in fused_patterns:
            self.assertGraphContainsExactly(graph, pat, 0)

    def findFusionGroups(self, graph):
        result = []
        for n in graph.nodes():
            if n.kind() == LLGA_FUSION_GROUP:
                result.append(n.g('Subgraph'))
                continue
            for block in n.blocks():
                result += self.findFusionGroups(block)
        return result

    def checkPatterns(self, graph, patterns):
        fusion_groups = self.findFusionGroups(graph)
        assert len(fusion_groups) == len(patterns), "length of subgraphs not equal to length of given patterns"

        for i in range(len(fusion_groups)):
            for pattern in patterns[i]:
                self.assertGraphContains(fusion_groups[i], pattern)

try:
    import torchvision
    HAS_TORCHVISION = True
except ImportError:
    HAS_TORCHVISION = False
except RuntimeError:
    HAS_TORCHVISION = False
skipIfNoTorchVision = unittest.skipIf(not HAS_TORCHVISION, 'no torchvision')

def get_eltwise_fn(name):
    if hasattr(torch, name):
        return getattr(torch, name)
    elif hasattr(F, name):
        return getattr(F, name)
    elif name == 'hardswish_':
        return torch.nn.Hardswish(inplace=True)
    else:
        raise NameError(f'Eltwise function {name} not found')


@unittest.skipIf(IS_AVX512_UNSUPPORTED, "This test fails for BF16 on machines without AVX512.")
@unittest.skipIf(LLGA_NOT_ENABLED, "MKL-DNN build is disabled")
class TestOp(JitLlgaTestCase):
    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_conv2d(self, dtype):
        for [spatial, in_channels, out_channels, kernel, padding, stride, dilation, g, bias] in itertools.product(
                [7, 8],
                [8, 15],
                [7, 16],
                [3, 4],
                [0, 2],
                [1, 2],
                [1, 2],
                [1, 2],
                [True, False]):

            m = nn.Conv2d(in_channels=in_channels * g,
                          out_channels=out_channels * g,
                          kernel_size=kernel,
                          padding=padding,
                          stride=stride,
                          dilation=dilation,
                          groups=g,
                          bias=bias)

            x = torch.rand(1, in_channels * g, spatial, spatial)
            _, graph = self.checkTrace(m, [x], dtype)
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_bn2d(self, dtype):
        m = nn.BatchNorm2d(32).eval()
        x = torch.rand(1, 32, 28, 28)
        _, graph = self.checkTrace(m, [x], dtype)
        # single-op partition shouldn't be created for softmax
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 0)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_eltwise(self, dtype):
        class M(nn.Module):
            def __init__(self, eltwise_fn):
                super().__init__()
                self.eltwise = eltwise_fn

            def forward(self, x):
                return self.eltwise(x)

        for eltwise in ['relu', 'gelu']:
            eltwise_fn = get_eltwise_fn(eltwise)
            m = M(eltwise_fn)
            x = torch.rand(1, 32, 28, 28)
            _, graph = self.checkTrace(m, [x], dtype)
            # single-op partition shouldn't be created.
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 0)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_max_pool2d(self, dtype):
        for [spatial, kernel, padding, stride, dilation, ceil_mode] in itertools.product(
                [15, 16, 17, 18, 19],
                [4, 5],
                [0, 1, 2],
                [1, 2],  # [1, 2, 4], TODO: fix issue in pad calculation
                [1],     # [1, 2], TODO: backend support for dilation
                [True, False]):

            m = nn.MaxPool2d(kernel_size=kernel,
                             stride=stride,
                             padding=padding,
                             dilation=dilation,
                             ceil_mode=ceil_mode)

            x = torch.rand(1, 4, spatial, spatial)
            _, graph = self.checkTrace(m, [x], dtype)
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_avg_pool2d(self, dtype):
        for [spatial, kernel, padding, stride, ceil_mode, count_include_pad] in itertools.product(
                [15, 16, 17, 18, 19],
                [4, 5],
                [0, 1, 2],
                [1, 2, 4],
                [False],  # TODO: oneDNN Graph does not fully support ceil_mode=True
                [True, False]):

            m = nn.AvgPool2d(kernel_size=kernel,
                             stride=stride,
                             padding=padding,
                             ceil_mode=ceil_mode,
                             count_include_pad=count_include_pad)

            x = torch.rand(1, 4, spatial, spatial)
            _, graph = self.checkTrace(m, [x], dtype)
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_variable_kernel_avg_pool2d(self, dtype):
        class M(nn.Module):
            def forward(self, x):
                x = F.avg_pool2d(x, kernel_size=(x.size(2), x.size(3)), padding=0, count_include_pad=False)
                return x

        x = torch.randn(1, 1000, 1, 1)
        m = M()
        _, graph = self.checkTrace(m, [x], dtype)
        # kernel_size is not Constant, shouldn't have any LLGA_FUSION_GROUP
        # TODO: with shape specialization, should have 1 LLGA_FUSION_GROUP
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 0)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_softmax(self, dtype):
        for dim in [-4, -3, -2, -1, 0, 1, 2, 3]:
            m = nn.Softmax(dim=dim)
            x = torch.rand(8, 12, 12, 12)
            _, graph = self.checkTrace(m, [x], dtype)
            # single-op partition shouldn't be created for softmax
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 0)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_linear(self, dtype):
        for bias in [True, False]:
            x = torch.rand(32, 28)
            m = torch.nn.Linear(in_features=28, out_features=64, bias=bias)
            _, graph = self.checkTrace(m, [x], dtype)
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
            self.assertFused(graph, ['aten::linear'])


    def _gen_binary_inputs(self, gen_permute=True):
        for xshape, yshape in [
            [[1, 32, 28, 28], [1, 32, 28, 28]],
            [[1, 32, 28, 28], [1, 1, 28, 28]],
            [[1, 32, 28, 28], [28]],
            [[1, 32, 28, 28], [1]],

        ]:
            yield torch.rand(xshape), torch.rand(yshape)
            if gen_permute and xshape != yshape:
                yield torch.rand(yshape), torch.rand(xshape)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_add(self, dtype):
        def forward_add(x, y):
            return torch.add(x, y, alpha=2)

        for x, y in self._gen_binary_inputs():
            _, graph = self.checkTrace(forward_add, [x, y], dtype)
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_add_scalar(self, dtype):
        def add_scalar(x):
            return 42 + x + 3.14

        x = torch.rand(32, 32)
        _, graph = self.checkTrace(add_scalar, [x], dtype)
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_addmm(self, dtype):
        # Just a sidenote - comparison of eager-mode & oneDNN Graph JIT outputs of
        # addmm (which entails matmul-bias-add fusion) might require higher tolerance
        # bounds for BF16. This is subject to change in the near future.
        def addmm(x, y, z):
            # alpha and beta are 1, by default
            return torch.addmm(z, x, y)

        x = torch.rand(64, 32)
        y = torch.rand(32, 32)
        z = torch.rand(64, 32)
        _, graph = self.checkTrace(addmm, [x, y, z], dtype)
        # single-op partition should be created for matmul with bias.
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_mul(self, dtype):
        def forward_mul(x, y):
            return torch.mul(x, y) * 3

        for x, y in self._gen_binary_inputs():
            _, graph = self.checkTrace(forward_mul, [x, y], dtype)
            # single-op partitions shouldn't be created
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_identity_binary(self, dtype):
        def forward(x):
            return x * 1 + 0.0

        x = torch.rand(32)
        _, graph = self.checkTrace(forward, [x], dtype)
        self.assertFused(graph, ['aten::add', 'aten::mul'])

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_layer_norm(self, dtype):
        # TODO: support more normalized_shape
        m = torch.nn.LayerNorm(10)
        x = torch.randn(2, 5, 10, 10)
        _, graph = self.checkTrace(m, [x], dtype)
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_cat(self, dtype):
        def cat_along_dim(d):
            def forward_cat(*inputs):
                return torch.cat(inputs, d)
            return forward_cat

        for xshape in [
            [8, 8, 8, 8],
            [64, 8, 32],
            [2048, 64],
        ]:
            for d in range(len(xshape)):
                x = torch.rand(xshape)
                _, graph = self.checkTrace(cat_along_dim(d), [x, x, x], dtype)
                self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_typecheck(self, dtype):
        x = torch.rand(32, 28, dtype=dtype)
        m = torch.nn.Linear(in_features=28, out_features=64, bias=True, dtype=dtype)
        traced, graph = self.checkTrace(m, [x], dtype)
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
        self.assertFused(graph, ['aten::linear'])
        # change the shape of the input, we should enter fallback graph
        x = torch.rand(5, 28, dtype=dtype)
        self.assertEqual(m(x), traced(x))


@unittest.skipIf(IS_AVX512_UNSUPPORTED, "This test fails for BF16 on machines without AVX512.")
@unittest.skipIf(LLGA_NOT_ENABLED, "MKL-DNN build is disabled")
class TestFusionPattern(JitLlgaTestCase):
    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_conv2d_eltwise(self, dtype):
        class M(nn.Module):
            def __init__(self, eltwise_fn):
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv2 = nn.Conv2d(32, 32, 3, padding=1, bias=False)
                self.eltwise = eltwise_fn

            def forward(self, x):
                x = self.conv1(x)
                x = self.eltwise(x)
                x = self.conv2(x)
                x = self.eltwise(x)
                return x

        for eltwise in ['relu', 'leaky_relu', 'sigmoid', 'square',
                        'abs', 'exp', 'hardswish', 'tanh', 'hardtanh']:
            for inplace in [True, False]:
                eltwise_fn_name = eltwise + '_' if inplace else eltwise
                eltwise_fn = get_eltwise_fn(eltwise_fn_name)

                m = M(eltwise_fn)
                x = torch.rand(1, 32, 28, 28)
                _, graph = self.checkTrace(m, [x], dtype=dtype)
                self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 2)
                # test if relu_ is replace with relu by mutation removal pass
                self.assertFused(graph, ['aten::' + eltwise_fn_name])
                # test if relu is fused into the fusion group
                self.assertFused(graph, ['aten::' + eltwise])

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_conv2d_silu(self, dtype):
        class M(nn.Module):
            def __init__(self, inplace):
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv2 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.eltwise = nn.SiLU(inplace=inplace)

            def forward(self, x):
                x = self.conv1(x)
                x = self.eltwise(x)
                x = self.conv2(x)
                return x
        for inplace in [False, True]:
            for memory_format in [torch.contiguous_format, torch.channels_last]:
                m = M(inplace)
                x = torch.rand(1, 32, 28, 28).to(memory_format=memory_format)

                _, graph = self.checkTrace(m, [x], dtype)
                self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 2)
                # oneDNN graph does not have silu OP. The bridge will convert silu to sigmoid - mul
                # Inplace op will become outplace op on the JIT graph
                patterns = [
                    ["aten::_convolution", 'aten::sigmoid', 'aten::mul'],
                    ["aten::_convolution"]
                ]
                silu_op = 'aten::silu_' if inplace else 'aten::silu'
                self.assertFused(graph, ['aten::_convolution', silu_op])
                self.checkPatterns(graph, patterns)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_ensure_tensor_is_rewrapped(self, dtype):
        class M(nn.Module):
            def __init__(self, eltwise_fn):
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv2 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv3 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv4 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.eltwise = eltwise_fn
                self.adaptive_avg_pool_2d = nn.AdaptiveAvgPool2d((5, 7))

            def forward(self, x, y):
                x = self.conv1(x)
                x = self.eltwise(x)
                x = self.conv2(x)
                x = self.eltwise(x)
                y = self.conv3(y)
                y = self.eltwise(y)
                y = self.conv4(y)
                y = self.eltwise(y)

                x = torch.add(x, y)
                x = self.adaptive_avg_pool_2d(x)
                return x

        eltwise_fn_name = 'relu'
        eltwise_fn = get_eltwise_fn(eltwise_fn_name)
        m = M(eltwise_fn)
        m = m.to(memory_format=torch.channels_last)
        x = torch.rand(1, 32, 28, 28).to(memory_format=torch.channels_last)
        y = torch.rand(1, 32, 28, 28).to(memory_format=torch.channels_last)
        # Simply test if the output is accurate
        # The output of the second partition is input to adaptive_avg_pool2d, which is
        # unsupported by LLGA. In resnext101 32x16d, we encountered an accuracy issue.
        _, graph = self.checkTrace(m, [x, y], dtype)
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 4)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_conv2d_clamp(self, dtype):
        class M(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv2 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv3 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv4 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv5 = nn.Conv2d(32, 32, 3, padding=1, bias=True)

            def forward(self, x):
                x = self.conv1(x)
                x = torch.clamp(x, min=float('-inf'))
                x = self.conv2(x)
                x = torch.clamp(x, min=-5)
                x = self.conv3(x)
                x = torch.clamp(x, min=0, max=float('inf'))
                x = self.conv4(x)
                x = torch.clamp(x, min=1, max=5)
                x = self.conv5(x)
                x = torch.clamp(x, max=2)
                return x

        for inplace in [False, True]:
            for memory_format in [torch.contiguous_format, torch.channels_last]:
                x = torch.rand(1, 32, 28, 28).to(memory_format=memory_format)
                m = M()
                _, graph = self.checkTrace(m, [x], dtype)
                self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 5)
                self.assertFused(graph, ['aten::_convolution', "aten::clamp"])

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_conv2d_bn(self, dtype):
        class M(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.bn1 = nn.BatchNorm2d(32)

            def forward(self, x):
                x = self.conv1(x)
                x = self.bn1(x)
                return x

        m = M().eval()
        if dtype == torch.bfloat16:
            m = optimization.fuse(m)
        x = torch.rand(1, 32, 28, 28)
        _, graph = self.checkTrace(m, [x], dtype)
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
        self.assertFused(graph, ['aten::_convolution', 'aten::batch_norm'])

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_conv2d_bn_relu(self, dtype):
        class M(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.bn1 = nn.BatchNorm2d(32)

            def forward(self, x):
                x = self.conv1(x)
                x = self.bn1(x)
                x = F.relu(x)
                return x

        m = M().eval()
        if dtype == torch.bfloat16:
            m = optimization.fuse(m)
        x = torch.rand(1, 32, 28, 28)
        _, graph = self.checkTrace(m, [x], dtype)
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
        self.assertFused(graph, ['aten::_convolution', 'aten::batch_norm',
                                 'aten::relu'])

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_bn2d_eltwise(self, dtype):
        class M(nn.Module):
            def __init__(self, eltwise_fn):
                super().__init__()
                self.eltwise = eltwise_fn
                self.bn = nn.BatchNorm2d(32)

            def forward(self, x):
                x = self.bn(x)
                x = self.eltwise(x)
                return x

        for eltwise in ['relu']:
            eltwise_fn = get_eltwise_fn(eltwise)
            m = M(eltwise_fn).eval()
            x = torch.rand(1, 32, 28, 28)
            _, graph = self.checkTrace(m, [x], dtype)
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
            self.assertFused(graph, ['aten::' + eltwise])

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_linear_eltwise(self, dtype):
        class M(nn.Module):
            def __init__(self, eltwise_fn, bias):
                super().__init__()
                self.linear = nn.Linear(28, 64, bias)
                self.eltwise = eltwise_fn

            def forward(self, x):
                x = self.linear(x)
                x = self.eltwise(x)
                return x

        for [has_bias, eltwise] in itertools.product(
                [True, False],
                ['relu', 'gelu', 'sigmoid', 'hardtanh', 'relu6', 'elu']):

            eltwise_fn = get_eltwise_fn(eltwise)
            m = M(eltwise_fn, has_bias)
            x = torch.rand(32, 28, requires_grad=False)
            _, graph = self.checkTrace(m, [x], dtype)
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
            self.assertFused(graph, ['aten::' + eltwise])

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_conv2d_sum(self, dtype):
        class M(nn.Module):
            def __init__(self, bias=False):
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=bias)
                self.bn1 = nn.BatchNorm2d(32)
                self.conv2 = nn.Conv2d(32, 32, 3, padding=1, bias=bias)
                self.bn2 = nn.BatchNorm2d(32)
                self.relu = nn.ReLU()
                self.conv3 = nn.Conv2d(32, 32, 3, padding=1, bias=bias)
                self.bn3 = nn.BatchNorm2d(32)

            def forward(self, x, y):
                x = self.conv1(x)
                x = self.bn1(x)
                y = self.conv2(y)
                y = self.bn2(y)
                z = self.relu(x + y)
                z = self.conv3(z)
                z = self.bn3(z)
                return z

        for bias in [True, False]:
            m = M(bias).eval()
            if dtype == torch.bfloat16:
                m = optimization.fuse(m)
            x = torch.rand(1, 32, 16, 16, requires_grad=False)
            y = torch.rand(1, 32, 16, 16, requires_grad=False)
            _, graph = self.checkTrace(m, [x, y], dtype)
            self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 3)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_wildcard(self, dtype):
        class M(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.eltwise = nn.ReLU()

            def forward(self, x):
                x = self.conv1(x)
                y = self.eltwise(x)
                return [x, y]

        # The pattern is as the following:
        #      conv
        #     |    \
        # eltwise   \
        #    |       \
        #  ListConstruct
        #
        # The output of conv is used by a wildcard op: ListConstruct.
        # Thus conv-eltwise cannot be selected into the same Partition.
        m = M()
        x = torch.rand(1, 32, 28, 28)
        _, graph = self.checkTrace(m, [x], dtype)
        # conv can exist in a single-op oneDNN Graph partition but not relu
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
        self.assertFused(graph, ['aten::_convolution'])

    @onlyCPU
    @dtypes(torch.int32)
    def test_wildcard_unsupported_dtype(self, dtype):
        class M(nn.Module):
            def forward(self, x):
                y = x // 2
                return y

        # In shufflenet_v2_x1_0, channels_per_groups is computed as:
        # channels_per_group = num_channels // groups
        # JIT IR converts groups to Long dtype, which is unsupported
        # by oneDNN Graph, viz. Long(requires_grad=0, device=cpu) = prim::Constant[value={2}]()
        # This test just ensures that the bridge code can handle
        # unsupported dtypes for inputs to ops unsupported
        # by oneDNN Graph. In this particular UT, aten::floor_divide
        # would be added as a wildcard in graph-construction stage.
        m = M()
        x = torch.tensor([32], dtype=dtype)
        _, graph = self.checkTrace(m, [x], dtype)
        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 0)

    @onlyCPU
    @dtypes(torch.float32, torch.bfloat16)
    def test_rewrap_tensor_input_to_pytorch(self, dtype):
        class M(nn.Module):
            def __init__(self, eltwise_fn):
                super().__init__()
                self.conv1 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.conv2 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
                self.eltwise = eltwise_fn
                self.adaptive_avg_pool_2d = nn.AdaptiveAvgPool2d((5, 7))

            def forward(self, x, y):
                x = self.conv1(x)
                x = self.eltwise(x)
                x = self.conv2(x)
                x = self.eltwise(x)
                x = torch.add(x, y)
                x = self.adaptive_avg_pool_2d(x)
                return x

        eltwise_fn_name = 'relu'
        eltwise_fn = get_eltwise_fn(eltwise_fn_name)
        m = M(eltwise_fn)
        m = m.to(memory_format=torch.channels_last)
        x = torch.rand(1, 32, 28, 28).to(memory_format=torch.channels_last)
        y = torch.rand(1, 32, 28, 28).to(memory_format=torch.channels_last)
        # Simply test if the output is accurate
        # The output of the second partition is input to adaptive_avg_pool2d, which is
        # unsupported by LLGA, so it must be handled by PyTorch, which should receive
        # correct strides info of the channels-last tensor.
        graph, _ = self.checkTrace(m, [x, y], dtype)

@unittest.skipIf(LLGA_NOT_ENABLED, "MKL-DNN build is disabled")
class TestEnableDisableLlgaFuser(JitTestCase):
    def setUp(self):
        super().setUp()
        self.is_enabled = torch._C._jit_set_llga_enabled(False)

    def tearDown(self):
        torch._C._jit_set_llga_enabled(self.is_enabled)
        super().tearDown()

    def test_context_manager(self):
        x = torch.randn(4, 8)
        y = torch.randn(4, 8)
        with torch.jit.fuser('fuser3'):
            with torch.jit.fuser('fuser3'):

                def t1(x, y):
                    o = x + y
                    o = o + 2.0
                    return o
                t_jit = torch.jit.script(t1)
                t_jit(x, y)
                t_jit(x, y)
                self.assertGraphContains(t_jit.graph_for(x, y), LLGA_FUSION_GROUP)

            def t2(x, y):
                o = x + y
                o = o + 3.0
                return o
            t_jit_2 = torch.jit.script(t2)
            t_jit_2(x, y)
            t_jit_2(x, y)
            self.assertGraphContains(t_jit_2.graph_for(x, y), LLGA_FUSION_GROUP)

        def t3(x, y):
            o = x + y
            o = o + 4.0
            return o
        t_jit_3 = torch.jit.script(t3)
        t_jit_3(x, y)
        t_jit_3(x, y)
        self.assertGraphContainsExactly(t_jit_3.graph_for(x, y), LLGA_FUSION_GROUP, 0)


@unittest.skipIf(LLGA_NOT_ENABLED, "MKL-DNN build is disabled")
@unittest.skip("Enable when integration with dynamo aot_autograd is more stable")
class TestDynamoAOT(JitTestCase):
    def test_dynamo_aot_ts_onednn(self):
        class Seq(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.layers = nn.Sequential(
                    nn.Linear(10, 10),
                    nn.ReLU(),
                    nn.Linear(10, 10),
                    nn.ReLU(),
                )

            def forward(self, x):
                return self.layers(x)

        mod = Seq()

        import torch._dynamo
        aot_mod = torch._dynamo.optimize("aot_ts", nopython=True)(mod)

        for _ in range(10):
            with torch.jit.fuser("fuser3"):
                loss = aot_mod(torch.rand([10, 10])).sum()
                loss.backward()

        torch._dynamo.reset()


@unittest.skipIf(IS_AVX512_UNSUPPORTED, "This test fails for BF16 on machines without AVX512.")
@unittest.skipIf(LLGA_NOT_ENABLED, "MKL-DNN build is disabled")
class TestModel(JitLlgaTestCase):
    @skipIfNoTorchVision
    def _test_vision(self, model_name, dtype):
        m = getattr(torchvision.models, model_name)().eval()
        if dtype == torch.bfloat16:
            m = optimization.fuse(m)
        x = torch.rand(1, 3, 224, 224) / 10
        _, graph = self.checkTrace(m, [x], dtype)
        self.assertFused(graph, ['aten::_convolution', 'aten::batch_norm',
                                 'aten::relu', 'aten::linear',
                                 'aten::avg_pool2d', 'aten::max_pool2d'])

for model_name, enabled in [
    ['resnet50', True],
    ['resnext50_32x4d', True],
    ['resnext101_32x8d', True],
    ['densenet121', True],
    ['densenet161', True],
    ['densenet169', True],
    ['densenet201', True],
    ['efficientnet_b0', True],
    ['efficientnet_b1', True],
    ['efficientnet_b2', True],
    ['efficientnet_b3', True],
    ['efficientnet_b4', True],
    ['efficientnet_b5', True],
    ['efficientnet_b6', True],
    ['efficientnet_b7', True],
    ['regnet_y_400mf', True],
    ['googlenet', TEST_SCIPY],
    ['mobilenet_v2', True],
    ['mobilenet_v3_large', True],
    ['mnasnet1_0', True],
    ['squeezenet1_0', True],
    ['vgg16', True],
    ['alexnet', True],
    ['shufflenet_v2_x1_0', True],
    ['wide_resnet50_2', True],
]:
    def _wrapper(mname, dtype):
        @unittest.skipIf(not enabled, 'Disabled')
        @separate_process
        def test(self, dtype=dtype):
            return self._test_vision(mname, dtype)
        return test

    for dtype in [torch.bfloat16, torch.float32]:
        setattr(TestModel, 'test_vision_{}_{}'.format(model_name, str(dtype).split("torch.")[1]), _wrapper(model_name, dtype))


instantiate_device_type_tests(TestFusionPattern, globals())
instantiate_device_type_tests(TestOp, globals())

if __name__ == '__main__':
    run_tests()