# Owner(s): ["module: inductor"]
import copy
import itertools
import os
import unittest
from typing import Callable, List, Optional

import torch
import torch._dynamo.config as dynamo_config
import torch._inductor.config as inductor_config
import torch._inductor.fx_passes.post_grad
import torch.nn.functional as F
from torch._dynamo.utils import count_calls, counters
from torch._higher_order_ops.auto_functionalize import auto_functionalized
from torch._higher_order_ops.out_dtype import out_dtype
from torch._inductor.fx_passes import joint_graph
from torch._inductor.pattern_matcher import (
    Arg,
    CallFunction,
    fwd_only,
    gen_pattern,
    is_mutation_op,
    KeywordArg,
    Match,
    PatternMatcherPass,
    PatternPrettyPrinter,
    register_graph_pattern,
    register_replacement,
    stable_topological_sort,
)
from torch._inductor.test_case import run_tests, TestCase
from torch._inductor.utils import run_and_get_code
from torch._inductor.virtualized import V
from torch.fx.experimental.proxy_tensor import make_fx
from torch.testing import FileCheck
from torch.testing._internal.common_cuda import SM80OrLater, xfailIfSM89
from torch.testing._internal.common_device_type import expectedFailureXPU, skipCUDAIf
from torch.testing._internal.common_utils import IS_LINUX, skipIfRocm, skipIfXpu
from torch.testing._internal.inductor_utils import (
    GPU_TYPE,
    HAS_GPU,
    IS_A100,
    IS_BIG_GPU,
)
from torch.utils import _pytree as pytree


aten = torch.ops.aten


class TestPatternMatcher(TestCase):
    device_type = GPU_TYPE

    def common(
        self,
        fn,
        args,
        expected_matches,
        expected_nodes,
        additional_check=lambda code: None,
        reference_in_float=False,
    ):
        counters.clear()
        torch.manual_seed(42)
        if reference_in_float:
            ref_inputs = pytree.tree_map_only(
                torch.Tensor, lambda x: x.to(torch.float32), args
            )
        else:
            ref_inputs = args
        expected = fn(*ref_inputs)
        torch.manual_seed(42)
        actual, codes = run_and_get_code(torch.compile(fn), *args)
        if len(codes) == 1:
            codes = codes[0]
        torch.testing.assert_close(actual, expected, check_dtype=not reference_in_float)

        self.assertEqual(
            counters["inductor"]["pattern_matcher_count"], expected_matches
        )
        self.assertEqual(counters["inductor"]["pattern_matcher_nodes"], expected_nodes)
        additional_check(codes)
        counters.clear()

    @inductor_config.patch(max_autotune_gemm=True)
    def test_mm_plus_mm(self):
        def fn(a, b, c, d):
            return torch.add(torch.mm(a, b), torch.mm(c, d))

        # when m1 == n1 and m2 == n2, mm_plus_mm can be matched to fused op
        fusible_args_list = [
            (
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
            ),
            (
                torch.randn(1, 4, device=GPU_TYPE),
                torch.randn(4, 2, device=GPU_TYPE),
                torch.randn(1, 5, device=GPU_TYPE),
                torch.randn(5, 2, device=GPU_TYPE),
            ),
        ]
        for args in fusible_args_list:
            self.common(fn, args, 1, 3)

        # if not fusible, it can only match add(mm())
        unfusible_args_list = [
            # https://github.com/pytorch/pytorch/issues/100670.
            (
                torch.randn(1, 4, device=GPU_TYPE),
                torch.randn(4, 2, device=GPU_TYPE),
                torch.randn(1, 2, device=GPU_TYPE),
                torch.randn(2, 1, device=GPU_TYPE),
            ),
            (
                torch.randn(1, 2, device=GPU_TYPE),
                torch.randn(2, 1, device=GPU_TYPE),
                torch.randn(1, 4, device=GPU_TYPE),
                torch.randn(4, 2, device=GPU_TYPE),
            ),
        ]
        for args in unfusible_args_list:
            self.common(fn, args, 1, 2)

    def _test_fused_int_mm_mul_impl(self, fn, args, fused_int_mm_mul_expected=True):
        torch._dynamo.reset()
        counters.clear()
        ref = fn(*args)
        test, (code,) = run_and_get_code(torch.compile(fn, mode="max-autotune"), *args)
        self.assertEqual("fused_int_mm_mul" in code, fused_int_mm_mul_expected)
        if fused_int_mm_mul_expected:
            indices = ~ref.isinf()
            torch.testing.assert_close(
                ref[indices], test[indices]
            )  # also checks that dtype is correct

    @skipIfRocm
    @skipIfXpu
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(force_fuse_int_mm_with_mul=True)
    def test_fused_int_mm_mul(self):
        def fn1(a, b, c):
            return out_dtype(torch.ops.aten.mm.default, torch.int32, a, b) * c

        def fn2(a, b, c):
            return (out_dtype(torch.ops.aten.mm.default, torch.int32, a, b) * c).to(
                torch.bfloat16
            )

        args_list = [
            (
                torch.randint(-128, 127, (32, 32), dtype=torch.int8, device=GPU_TYPE),
                torch.randint(-128, 127, (32, 8), dtype=torch.int8, device=GPU_TYPE),
                torch.randn((32, 1), dtype=torch.float16, device=GPU_TYPE) * 0 + 0.5,
            ),
            (
                torch.randint(-128, 127, (32, 32), dtype=torch.int8, device=GPU_TYPE),
                torch.randint(-128, 127, (32, 8), dtype=torch.int8, device=GPU_TYPE),
                torch.randn((1, 8), dtype=torch.bfloat16, device=GPU_TYPE),
            ),
            (
                torch.randint(-128, 127, (32, 32), dtype=torch.int8, device=GPU_TYPE),
                torch.randint(-128, 127, (32, 8), dtype=torch.int8, device=GPU_TYPE),
                torch.randn((1, 8), dtype=torch.float32, device=GPU_TYPE),
            ),
        ]

        for args in args_list:
            self._test_fused_int_mm_mul_impl(fn1, args, True)
            self._test_fused_int_mm_mul_impl(fn2, args, True)

    @skipIfRocm
    @skipIfXpu
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(force_fuse_int_mm_with_mul=True)
    def test_fused_int_mm_mul_epilogue(self):
        def fn1(a, b, c):
            return (
                (out_dtype(torch.ops.aten.mm.default, torch.int32, a, b) * c) * 0.5
            ).relu()

        def fn2(a, b, c):
            return (
                (out_dtype(torch.ops.aten.mm.default, torch.int32, a, b) * c).to(
                    torch.bfloat16
                )
                * 0.5
            ).relu()

        args_list = [
            (
                torch.randint(-128, 127, (32, 32), dtype=torch.int8, device=GPU_TYPE),
                torch.randint(-128, 127, (32, 8), dtype=torch.int8, device=GPU_TYPE),
                torch.randn((32, 1), dtype=torch.float16, device=GPU_TYPE) * 0 + 0.5,
            ),
            (
                torch.randint(-128, 127, (32, 32), dtype=torch.int8, device=GPU_TYPE),
                torch.randint(-128, 127, (32, 8), dtype=torch.int8, device=GPU_TYPE),
                torch.randn((1, 8), dtype=torch.bfloat16, device=GPU_TYPE),
            ),
            (
                torch.randint(-128, 127, (32, 32), dtype=torch.int8, device=GPU_TYPE),
                torch.randint(-128, 127, (32, 8), dtype=torch.int8, device=GPU_TYPE),
                torch.randn((1, 8), dtype=torch.float32, device=GPU_TYPE),
            ),
        ]

        for args in args_list:
            self._test_fused_int_mm_mul_impl(fn1, args, True)
            self._test_fused_int_mm_mul_impl(fn2, args, True)

    @skipIfRocm
    @skipIfXpu
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(force_fuse_int_mm_with_mul=True)
    def test_fused_int_mm_mul_gating(self):
        def fn1(a, b, c):
            return out_dtype(torch.ops.aten.mm.default, torch.int32, a, b) * c

        args1 = (
            torch.randint(-128, 127, (32, 32), dtype=torch.int8, device=GPU_TYPE),
            torch.randint(-128, 127, (32, 8), dtype=torch.int8, device=GPU_TYPE),
            torch.randn((8), dtype=torch.float32, device=GPU_TYPE),
        )

        args2 = (
            torch.randint(-128, 127, (32, 32), dtype=torch.int8, device=GPU_TYPE),
            torch.randint(-128, 127, (32, 8), dtype=torch.int8, device=GPU_TYPE),
            torch.randn((32, 1), dtype=torch.float16, device=GPU_TYPE),
        )
        self._test_fused_int_mm_mul_impl(fn1, args1, False)
        self._test_fused_int_mm_mul_impl(fn1, [arg.cpu() for arg in args2], False)
        inductor_config.force_fuse_int_mm_with_mul = False
        self._test_fused_int_mm_mul_impl(fn1, args2, False)

    def _test_mixed_impl(
        self,
        fn,
        args,
        mixed_mm_expected,
        fallback_mixed_mm_expected,
        rtol=None,
        atol=None,
    ):
        torch._dynamo.reset()
        counters.clear()
        ref = fn(*args)
        test, (code,) = run_and_get_code(torch.compile(fn), *args)
        torch.testing.assert_close(ref, test, rtol=rtol, atol=atol)
        self.assertEqual("mixed_mm" in code, mixed_mm_expected)
        self.assertEqual("fallback_mixed_mm" in code, fallback_mixed_mm_expected)

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(mixed_mm_choice="triton")
    def test_mixed_mm(self):
        def fn(a, b):
            return torch.mm(a, b.to(a.dtype))

        args_list = [
            (
                torch.randn(8, 8, device=GPU_TYPE),
                torch.randint(-128, 127, (8, 8), dtype=torch.int8, device=GPU_TYPE),
            ),
            (
                torch.randn(8, 2, device=GPU_TYPE, dtype=torch.bfloat16),
                torch.randint(-128, 127, (2, 8), dtype=torch.int8, device=GPU_TYPE),
            ),
            (
                torch.randn(8, 5, device=GPU_TYPE, dtype=torch.float16),
                torch.randint(0, 255, (5, 2), dtype=torch.uint8, device=GPU_TYPE),
            ),
            (
                torch.randn(8, 8, device=GPU_TYPE, dtype=torch.float32),
                torch.randn(8, 8, device=GPU_TYPE, dtype=torch.bfloat16),
            ),
        ]

        for args in args_list:
            self._test_mixed_impl(fn, args, True, False)

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(mixed_mm_choice="triton")
    def test_mixed_mm_exhaustive_dtypes(self):
        def fn(a, b):
            return torch.mm(a, b.to(a.dtype))

        dtypes_left = [torch.float16, torch.float32, torch.bfloat16]
        dtypes_right = [torch.int8, torch.uint8]
        dtype_ranges = {torch.uint8: (0, 255), torch.int8: (-128, 127)}
        for dtype_left, dtype_right in itertools.product(dtypes_left, dtypes_right):
            low, high = dtype_ranges[dtype_right]
            args = (
                torch.randn(256, 256, dtype=dtype_left, device=GPU_TYPE),
                torch.randint(
                    low, high, (256, 256), dtype=dtype_right, device=GPU_TYPE
                ),
            )
            fallback_mixed_mm_expected = (
                dtype_left == torch.bfloat16 and dtype_right == torch.uint8
            )
            self._test_mixed_impl(
                fn, args, True, fallback_mixed_mm_expected, rtol=0.16, atol=1e-4
            )

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(mixed_mm_choice="triton")
    def test_mixed_mm_bad_cases(self):
        def fn(a, b):
            return torch.mm(a, b.to(a.dtype))

        # when b is transposed and not contiguous, we skip triton and use fallback
        args_list = [
            (
                torch.randn(8, 8, device=GPU_TYPE, dtype=torch.float16),
                torch.randint(-128, 127, (4, 8), dtype=torch.int8, device=GPU_TYPE).t()[
                    :, ::2
                ],
            ),
            (
                torch.randn(8, 8, device=GPU_TYPE, dtype=torch.bfloat16),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device=GPU_TYPE).t()[
                    :, ::2
                ],
            ),
        ]

        for args in args_list:
            self._test_mixed_impl(fn, args, True, True)

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(mixed_mm_choice="triton", max_autotune_gemm=True)
    def test_mixed_mm_epi_works(self):
        def fn(a, b, c, d):
            return torch.mm(a, b.to(a.dtype)) * c + d

        args_list = [
            (
                torch.randn(8, 8, device=GPU_TYPE),
                torch.randint(-128, 127, (8, 8), dtype=torch.int8, device=GPU_TYPE),
                torch.randn(8, device=GPU_TYPE),
                torch.randn(8, device=GPU_TYPE),
            ),
            (
                torch.randn(8, 2, device=GPU_TYPE, dtype=torch.bfloat16),
                torch.randint(-128, 127, (2, 8), dtype=torch.int8, device=GPU_TYPE),
                torch.randn(8, device=GPU_TYPE, dtype=torch.bfloat16),
                torch.randn(8, device=GPU_TYPE, dtype=torch.bfloat16),
            ),
            (
                torch.randn(8, 5, device=GPU_TYPE, dtype=torch.float16),
                torch.randint(0, 255, (5, 2), dtype=torch.uint8, device=GPU_TYPE),
                torch.randn(2, device=GPU_TYPE, dtype=torch.float16),
                torch.randn(2, device=GPU_TYPE, dtype=torch.float16),
            ),
        ]

        for args in args_list:
            self._test_mixed_impl(fn, args, True, False)

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @skipCUDAIf(not IS_A100, "heuristic only run on Linux A100")
    @skipCUDAIf(not IS_BIG_GPU, "tests fail on small GPU")
    @inductor_config.patch(
        mixed_mm_choice="heuristic",
        autoheuristic_use="",
        fx_graph_cache=False,
        fx_graph_remote_cache=False,
        shape_padding=False,
    )
    def test_mixed_mm_heuristic_no(self):
        def fn(a, b):
            return torch.mm(a, b.to(a.dtype))

        # examples that should not be selected by handwritten heuristic
        mat1_dtype = torch.float16
        dyn_tensor = torch.randn(4, 4096, dtype=mat1_dtype, device=GPU_TYPE)
        torch._dynamo.mark_dynamic(dyn_tensor, 0)
        args_list = [
            (
                torch.randn(1, 4097, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4097, 4096), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
            (
                torch.randn(1, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 4097), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
            (
                torch.randn(8, 8, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(-128, 127, (8, 8), dtype=torch.int8, device=GPU_TYPE),
            ),
            (
                torch.randn(8, 2048, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (2048, 2048), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
            (
                torch.randn(8, 2048, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (2048, 2048), dtype=torch.int8, device=GPU_TYPE
                ).t(),
            ),
            (
                torch.randn(8, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 4096), dtype=torch.int8, device=GPU_TYPE
                )[:, ::2],
            ),
            (
                torch.randn(1, 4096, dtype=torch.float32, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 4096), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
            (
                dyn_tensor,
                torch.randint(
                    -128, 127, (4096, 4096), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
        ]

        for args in args_list:
            self._test_mixed_impl(fn, args, True, True)

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @skipCUDAIf(not IS_A100, "heuristic only run on Linux A100")
    @skipCUDAIf(not IS_BIG_GPU, "tests fail on small GPU")
    @inductor_config.patch(
        mixed_mm_choice="heuristic",
        autoheuristic_use="",
        fx_graph_cache=False,
        fx_graph_remote_cache=False,
        shape_padding=False,
    )
    def test_mixed_mm_heuristic_yes(self):
        def fn(a, b):
            return torch.mm(a, b.to(a.dtype))

        mat1_dtype = torch.float16
        # examples that should be selected by handwritten heuristic
        args_list = [
            (
                torch.randn(1, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 4096), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
            (
                torch.randn(4, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 4096), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
            (
                torch.randn(8, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 4096), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
            (
                torch.randn(8, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 4096), dtype=torch.int8, device=GPU_TYPE
                ).t(),
            ),
            (
                torch.randn(16, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (8192, 4096), dtype=torch.int8, device=GPU_TYPE
                ).t(),
            ),
            (
                torch.randn(32, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 8192), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
            (
                torch.randn(64, 4096, dtype=mat1_dtype, device=GPU_TYPE),
                torch.randint(
                    -128, 127, (4096, 4096), dtype=torch.int8, device=GPU_TYPE
                ),
            ),
        ]

        for args in args_list:
            self._test_mixed_impl(fn, args, True, False, rtol=0.01, atol=0.04)

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    def test_mixed_mm_gating(self):
        def fn(a, b):
            return torch.mm(a, b.to(a.dtype))

        args = (
            torch.randn(8, 8, device=GPU_TYPE),
            torch.randint(-128, 127, (8, 8), dtype=torch.int8, device=GPU_TYPE),
        )
        # will ignore the mixed_mm code (including fallback)
        with inductor_config.patch(
            {"mixed_mm_choice": "default", "use_mixed_mm": False}
        ):
            self._test_mixed_impl(fn, args, False, False)

        # will use fallback_mixed_mm kernel due to no gemm_autotune
        with inductor_config.patch(
            {"mixed_mm_choice": "default", "use_mixed_mm": True}
        ):
            self._test_mixed_impl(fn, args, True, True)

        # will use mixed_mm kernel
        with inductor_config.patch(
            {"mixed_mm_choice": "triton", "use_mixed_mm": False}
        ):
            self._test_mixed_impl(fn, args, True, False)

        # shows that use_mixed_mm doesn't do anything if foce_mixed_mm is set
        with inductor_config.patch({"mixed_mm_choice": "triton", "use_mixed_mm": True}):
            self._test_mixed_impl(fn, args, True, False)

        # will use fallback_mixed_mm kernel
        with inductor_config.patch({"mixed_mm_choice": "aten", "use_mixed_mm": False}):
            self._test_mixed_impl(fn, args, True, True)

        # will use fallback_mixed_mm kernel
        with inductor_config.patch({"mixed_mm_choice": "aten", "use_mixed_mm": True}):
            self._test_mixed_impl(fn, args, True, True)

        # will use fallback_mixed_mm kernel because fallback is the only choice
        with inductor_config.patch(
            {"mixed_mm_choice": "aten", "use_mixed_mm": True, "max_autotune_gemm": True}
        ):
            self._test_mixed_impl(fn, args, True, True)

    @inductor_config.patch(use_mixed_mm=True)
    def test_mixed_mm_cpu(self):
        def fn(a, b):
            return torch.mm(a, b.to(a.dtype))

        args = (
            torch.randn(8, 8),
            torch.randint(-128, 127, (8, 8), dtype=torch.int8),
        )
        self._test_mixed_impl(fn, args, False, False)

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(use_mixed_mm=True)
    def test_uint4x2_mixed_mm(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        def check_uint4x2_mixed_mm(args, expect_mixed_mm):
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            torch.testing.assert_close(ref, test)
            self.assertEqual("uint4x2_mixed_mm" in code, expect_mixed_mm)

        args_expect_mixed_mm = [
            (
                torch.randn(8, 8, device=GPU_TYPE),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device=GPU_TYPE),
            ),
            (
                torch.randn(8, 8, device=GPU_TYPE, dtype=torch.float16),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device=GPU_TYPE)
                .t()
                .contiguous()
                .t(),
            ),
        ]

        for args in args_expect_mixed_mm:
            check_uint4x2_mixed_mm(args, True)

        # mixed mm is only enabled when casting from a lower-bitwidth dtype to a higher one
        args_expect_no_mixed_mm = [
            (
                torch.randn(8, 8, device=GPU_TYPE),
                torch.randint(0, 255, (4, 8), dtype=torch.int32, device=GPU_TYPE),
            ),
            (
                torch.randn(8, 8, device=GPU_TYPE),
                torch.randint(0, 255, (4, 8), dtype=torch.int64, device=GPU_TYPE),
            ),
        ]

        for args in args_expect_no_mixed_mm:
            check_uint4x2_mixed_mm(args, False)

    @expectedFailureXPU
    @skipCUDAIf(not SM80OrLater, "need sm_80")
    @inductor_config.patch(use_mixed_mm=True)
    def test_uint4x2_mixed_mm_epi(self):
        def fn(a, b, c, d):
            return (
                torch.mm(
                    a,
                    torch.cat((b & 0xF, b >> 4), 1)
                    .reshape(-1, b.shape[1])
                    .to(a.dtype)
                    .sub(8),
                )
                * c
                + d
            )

        args_list = [
            (
                torch.randn(8, 8, device=GPU_TYPE),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device=GPU_TYPE),
                torch.randn(8, device=GPU_TYPE),
                torch.randn(8, device=GPU_TYPE),
            ),
        ]

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            torch.testing.assert_close(ref, test)
            self.assertTrue("uint4x2_mixed_mm" in code)
            self.assertTrue("fused_add_mm_mul" in code)

    @inductor_config.patch(use_mixed_mm=True)
    def test_uint4x2_mixed_mm_fail_to_match(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        args_list = [
            (  # cpu
                torch.randn(8, 8),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8),
            ),
            (  # int8
                torch.randn(8, 8, device=GPU_TYPE),
                torch.randint(-128, 127, (4, 8), dtype=torch.int8, device=GPU_TYPE),
            ),  # we don't match for int8 since numerics
        ]  # for int8 bitshifts don't match between triton and pytorch

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            torch.testing.assert_close(ref, test)
            self.assertFalse("uint4x2_mixed_mm" in code)

    @inductor_config.patch(mixed_mm_choice="default")
    @inductor_config.patch(use_mixed_mm=False)
    def test_uint4x2_mixed_mm_gating_works(self):
        def fn(a, b):
            return torch.mm(
                a,
                torch.cat((b & 0xF, b >> 4), 1)
                .reshape(-1, b.shape[1])
                .to(a.dtype)
                .sub(8),
            )

        args_list = [
            (
                torch.randn(8, 8, device=GPU_TYPE),
                torch.randint(0, 255, (4, 8), dtype=torch.uint8, device=GPU_TYPE),
            ),
        ]

        for args in args_list:
            torch._dynamo.reset()
            counters.clear()
            ref = fn(*args)
            test, (code,) = run_and_get_code(torch.compile(fn), *args)
            torch.testing.assert_close(ref, test)
            self.assertFalse("uint4x2_mixed_mm" in code)

    def test_addmm(self):
        def fn(a, b, c):
            return torch.add(a, torch.mm(b, c)), torch.mm(b, c) + a

        args_list = [
            (
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                True,
            ),
            (
                torch.randn(8, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(16, 8, device=GPU_TYPE),
                True,
            ),
            (
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(1, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                False,
            ),
            (
                torch.randn(1, 16, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                False,
            ),
            (
                4,
                torch.randn(16, 16, device=GPU_TYPE),
                torch.randn(16, 16, device=GPU_TYPE),
                False,
            ),
        ]
        for a, b, c, should_fuse in args_list:
            torch._dynamo.reset()
            counters.clear()
            args = (a, b, c)
            e1, e2 = fn(*args)
            a1, a2 = torch.compile(fn)(*args)
            torch.testing.assert_close(a1, e1)
            torch.testing.assert_close(a2, e2)
            count, nodes = (2, 4) if should_fuse else (0, 0)
            self.assertEqual(counters["inductor"]["pattern_matcher_count"], count)
            self.assertEqual(counters["inductor"]["pattern_matcher_nodes"], nodes)

    def test_addmm_symbolic_scalar(self):
        def fn(m1, m2):
            bias = m1.size(0)
            return torch.add(bias, torch.mm(m1, m2)), torch.mm(m1, m2) + bias

        m1 = torch.randn(16, 16, device=GPU_TYPE)
        m2 = torch.randn(16, 16, device=GPU_TYPE)

        counters.clear()
        expect = fn(m1, m2)
        actual = torch.compile(fn, dynamic=True)(m1, m2)
        self.assertEqual(expect, actual)
        self.assertEqual(counters["inductor"]["pattern_matcher_count"], 0)

    def test_addmm_broadcasting_bias(self):
        class Model(torch.nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.linear = torch.nn.functional.linear
                self.linear_weight = torch.randn(4, 4).to(GPU_TYPE)
                self.bias = torch.randn(1, 4).to(GPU_TYPE)

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

        input_tensor = torch.randn(1, 3, 4).to(GPU_TYPE)

        func = Model().to(GPU_TYPE)

        res1 = func(input_tensor)
        jit_func = torch.compile(func)
        res2 = jit_func(input_tensor)

        self.assertEqual(res1, res2)

    def test_cat_mm(self):
        def fn(a, b, c):
            return torch.cat(
                [
                    torch.mm(a, b),
                    torch.mm(b, c),
                    torch.mm(a, c),
                ],
                1,
            )

        args = [
            torch.randn(16, 16, device=GPU_TYPE),
            torch.randn(16, 16, device=GPU_TYPE),
            torch.randn(16, 16, device=GPU_TYPE),
        ]
        out, code = run_and_get_code(torch.compile(fn), *args)
        self.assertEqual(out, fn(*args))
        FileCheck().check("call").check_not(".run").run(code[0])

    def test_cat_addmm(self):
        def fn(a, b, c):
            return torch.cat(
                [
                    torch.addmm(a, b, c),
                    torch.addmm(b, c, a),
                    torch.addmm(c, a, b),
                ],
                1,
            )

        args = [
            torch.randn(16, 16, device=GPU_TYPE),
            torch.randn(16, 16, device=GPU_TYPE),
            torch.randn(16, 16, device=GPU_TYPE),
        ]
        out, code = run_and_get_code(torch.compile(fn), *args)
        self.assertEqual(out, fn(*args))
        FileCheck().check("call").check_not(".run").run(code[0])

    def test_cat_slice_cat_cuda(self):
        def fn(a, b):
            cat_1 = torch.ops.aten.cat.default([a, b], 1)
            slice_1 = torch.ops.aten.slice.Tensor(cat_1, 0, 0, 9223372036854775807)
            slice_2 = torch.ops.aten.slice.Tensor(slice_1, 1, 0, 19)
            return torch.ops.aten.cat.default([cat_1, slice_2], 1)

        args = [
            torch.randn(2, 32, device=GPU_TYPE),
            torch.randn(2, 16, device=GPU_TYPE),
        ]
        self.common(fn, args, 1, 3)

        args = [
            torch.randn(2, 8, device=GPU_TYPE),
            torch.randn(2, 16, device=GPU_TYPE),
        ]
        torch._dynamo.reset()
        counters.clear()
        expected = fn(*args)
        actual = torch.compile(fn)(*args)
        torch.testing.assert_close(actual, expected)
        # We don't recompile for dynamic-shape cases.
        if dynamo_config.assume_static_by_default:
            self.assertEqual(counters["inductor"]["pattern_matcher_count"], 1)
            self.assertEqual(counters["inductor"]["pattern_matcher_nodes"], 3)

        # Verify we fallback to non-optimal path for negative `end`.
        def fn(a, b):
            cat_1 = torch.ops.aten.cat.default([a, b], 1)
            slice_1 = torch.ops.aten.slice.Tensor(cat_1, 0, 0, 9223372036854775807)
            slice_2 = torch.ops.aten.slice.Tensor(slice_1, 1, 0, -1)
            return torch.ops.aten.cat.default([cat_1, slice_2], 1)

        args = [
            torch.randn(2, 8, device=GPU_TYPE),
            torch.randn(2, 16, device=GPU_TYPE),
        ]
        self.common(fn, args, 1, 3)

    def test_pointless_view_pair(self):
        def f(x):
            x = aten.view.default(x, [3, 5, 7])
            x = aten.view.default(x, [15, 7])
            return x

        x = torch.randn(15, 7, device=GPU_TYPE)
        gm = make_fx(f)(x)
        self.assertEqual(count_calls(gm.graph), 2)
        joint_graph.joint_graph_passes(gm)
        self.assertEqual(count_calls(gm.graph), 0)

        def f(x):
            x1 = aten.view.default(x, [3, 5, 7])
            x2 = aten.view.default(x1, [15, 7])
            return x1, x2

        gm = make_fx(f)(x)
        self.assertEqual(count_calls(gm.graph), 2)
        joint_graph.joint_graph_passes(gm)
        self.assertEqual(count_calls(gm.graph), 2)

    def test_pointless_permute_pair(self):
        def f(x):
            x = aten.permute.default(x, [1, 0])
            x = aten.permute.default(x, [1, 0])
            return x

        x = torch.randn(15, 7, device=GPU_TYPE)
        gm = make_fx(f)(x)
        self.assertEqual(count_calls(gm.graph), 2)
        joint_graph.joint_graph_passes(gm)
        self.assertEqual(count_calls(gm.graph), 0)

        def f(x):
            x1 = aten.permute.default(x, [1, 0])
            x2 = aten.permute.default(x1, [1, 0])
            return x1, x2

        gm = make_fx(f)(x)
        self.assertEqual(count_calls(gm.graph), 2)
        joint_graph.joint_graph_passes(gm)
        self.assertEqual(count_calls(gm.graph), 2)

    def test_pointless_permute_pair_3d(self):
        def f(x):
            x = aten.permute.default(x, [1, 0, 2])
            x = aten.permute.default(x, [1, 0, 2])
            return x

        x = torch.randn(3, 5, 7, device=GPU_TYPE)
        gm = make_fx(f)(x)
        self.assertEqual(count_calls(gm.graph), 2)
        joint_graph.joint_graph_passes(gm)
        self.assertEqual(count_calls(gm.graph), 0)

        def f(x):
            x1 = aten.permute.default(x, [1, 0, 2])
            x2 = aten.permute.default(x1, [1, 0, 2])
            return x1, x2

        gm = make_fx(f)(x)
        self.assertEqual(count_calls(gm.graph), 2)
        joint_graph.joint_graph_passes(gm)
        self.assertEqual(count_calls(gm.graph), 2)

    def test_pointless_convert(self):
        def fn1(x):
            x = torch.ops.prims.convert_element_type.default(x, torch.float16)
            x = torch.ops.prims.convert_element_type.default(x, torch.float32)
            return x

        gm = torch.fx.symbolic_trace(fn1)
        self.assertEqual(count_calls(gm.graph), 2)
        joint_graph.joint_graph_passes(gm)
        self.assertEqual(count_calls(gm.graph), 1)

        def fn2(x):
            x = torch.ops.prims.convert_element_type.default(x, torch.int32)
            x = torch.ops.prims.convert_element_type.default(x, torch.float32)
            return x

        gm = torch.fx.symbolic_trace(fn2)
        self.assertEqual(count_calls(gm.graph), 2)
        joint_graph.joint_graph_passes(gm)
        self.assertEqual(count_calls(gm.graph), 2)

    # Constant folding was explicitly turned off due to issue #108388
    # Turn it back on for test
    @inductor_config.patch(joint_graph_constant_folding=True)
    def test_pointless_cumsum(self):
        def fn1():
            ones = torch.full(
                [1, 128], 1, layout=torch.strided, dtype=torch.float32
            ).to(torch.int64)
            return torch.cumsum(ones, 1) * ones

        def fn2():
            ones = torch.full(
                [55, 10], 1, layout=torch.strided, dtype=torch.float32
            ).to(torch.int64)
            return torch.cumsum(ones, 1)

        def fn3():
            twos = torch.full([5, 4, 3], 2, dtype=torch.int64)
            return torch.cumsum(twos, 0)

        def fn4():
            x = torch.full([100], 0.1, dtype=torch.float32)
            return torch.cumsum(x, 0)

        def fn5():
            t1 = torch.full([2, 4], 1)
            t2 = t1.to(dtype=torch.bool)
            return torch.cumsum(t2, 1)

        def fn6():
            x = torch.full([10, 10], True, dtype=torch.int32)
            return torch.cumsum(x, 1)

        for fn in (fn1, fn2, fn3, fn4, fn5, fn6):
            result, (code,) = run_and_get_code(torch.compile(fn, fullgraph=True))
            self.assertNotIn("aten.cumsum", code)
            self.assertEqual(result, fn())
            self.assertEqual(counters["inductor"]["pattern_matcher_count"], 1)
            counters.clear()

    def test_splitwithsizes_cat(self):
        # Good case
        def fn(a):
            split_with_sizes = torch.ops.aten.split_with_sizes.default(a, [8, 24], 1)
            getitem = split_with_sizes[0]
            getitem_1 = split_with_sizes[1]
            cat = torch.ops.aten.cat.default([getitem, getitem_1], 1)
            return cat**2

        args = [
            torch.randn(2, 32, device=GPU_TYPE),
        ]
        self.common(fn, args, 1, 4)

        # Not all getitems are passed to cat
        def fn(a):
            split_with_sizes = torch.ops.aten.split_with_sizes.default(a, [8, 8, 16], 1)
            getitem = split_with_sizes[0]
            getitem_1 = split_with_sizes[1]
            getitem_2 = split_with_sizes[2]
            cat = torch.ops.aten.cat.default([getitem, getitem_1], 1)
            return cat**2 + getitem_2

        args = [
            torch.randn(2, 32, device=GPU_TYPE),
        ]
        self.common(fn, args, 0, 0)

        # Different dimensions  (TODO this case should be handled by replacing with a reshape)
        def fn(a):
            split_with_sizes = torch.ops.aten.split_with_sizes.default(
                a, [8, 8, 8, 8], 1
            )
            cat = torch.ops.aten.cat.default(split_with_sizes, 0)
            return cat**2

        args = [
            torch.randn(2, 32, device=GPU_TYPE),
        ]
        self.common(fn, args, 0, 0)

        # https://github.com/pytorch/pytorch/issues/99686.
        def fn(a):
            x = torch.ops.aten.split_with_sizes.default(a, [3, 2, 3], dim=1)
            cat = torch.ops.aten.cat.default([x[1], x[0], x[2]], dim=1)
            return cat

        args = [
            torch.randn(1, 8, device=GPU_TYPE),
        ]
        self.common(fn, args, 0, 0)

    def test_cat_splitwithsizes(self):
        # good case
        def fn(a, b, c):
            cat = torch.ops.aten.cat.default([a, b, c], 1)
            split_with_sizes = torch.ops.aten.split_with_sizes.default(
                cat, [2, 3, 5], 1
            )
            return [s**2 for s in split_with_sizes]

        args = [
            torch.randn(2, 2, device=GPU_TYPE),
            torch.randn(2, 3, device=GPU_TYPE),
            torch.randn(2, 5, device=GPU_TYPE),
        ]
        self.common(fn, args, 1, 2)

        # cat node has other users
        def fn(a, b, c):
            cat = torch.ops.aten.cat.default([a, b, c], 1)
            split_with_sizes = torch.ops.aten.split_with_sizes.default(
                cat, [2, 3, 5], 1
            )
            return [s**2 for s in split_with_sizes] + [cat**3]

        args = [
            torch.randn(2, 2, device=GPU_TYPE),
            torch.randn(2, 3, device=GPU_TYPE),
            torch.randn(2, 5, device=GPU_TYPE),
        ]
        self.common(fn, args, 0, 0)

        # cat and split dims are different
        def fn(a, b, c):
            cat = torch.ops.aten.cat.default([a, b, c], 1)
            split_with_sizes = torch.ops.aten.split_with_sizes.default(
                cat, [2, 3, 5], 0
            )
            return [s**2 for s in split_with_sizes]

        args = [
            torch.randn(10, 2, device=GPU_TYPE),
            torch.randn(10, 3, device=GPU_TYPE),
            torch.randn(10, 5, device=GPU_TYPE),
        ]
        self.common(fn, args, 0, 0)

        # cat and split lenghts are different
        def fn(a, b, c):
            cat = torch.ops.aten.cat.default([a, b, c], 1)
            split_with_sizes = torch.ops.aten.split_with_sizes.default(cat, [5, 5], 1)
            return [s**2 for s in split_with_sizes]

        args = [
            torch.randn(2, 2, device=GPU_TYPE),
            torch.randn(2, 3, device=GPU_TYPE),
            torch.randn(2, 5, device=GPU_TYPE),
        ]
        self.common(fn, args, 0, 0)

        # cat input sizes and split sizes are different
        def fn(a, b, c):
            cat = torch.ops.aten.cat.default([a, b, c], 1)
            split_with_sizes = torch.ops.aten.split_with_sizes.default(
                cat, [2, 5, 3], 1
            )
            return [s**2 for s in split_with_sizes]

        args = [
            torch.randn(2, 2, device=GPU_TYPE),
            torch.randn(2, 3, device=GPU_TYPE),
            torch.randn(2, 5, device=GPU_TYPE),
        ]
        self.common(fn, args, 0, 0)

    def test_symint_pattern_matching(self):
        import torch._inductor.config as config
        from torch._inductor.pattern_matcher import (
            fwd_only,
            PatternMatcherPass,
            register_replacement,
        )

        saved_graph = None

        class _CustomPass(PatternMatcherPass):
            def __init__(self) -> None:
                super().__init__()

            def __call__(self, g: torch.fx.graph.Graph):
                self.apply(g)
                nonlocal saved_graph
                saved_graph = g

        with config.patch(
            # leave custom pass only in post_grad_passes()
            pattern_matcher=False,
            # define pattern match as custom post grad opt pass
            post_grad_custom_pre_pass=None,
            post_grad_custom_post_pass=_CustomPass(),
        ):

            def add(x, y):
                return x + y

            # testing that
            def sym_minus(x, y):
                return (x - (-y.size(0))) - (y * -1) - y.size(0)

            device = "cpu"
            my_args = [
                torch.empty([8, 1], device=device),
                torch.empty([10], device=device),
            ]

            invoked = False

            def extra_check(match):
                nonlocal invoked
                invoked = True
                return True

            register_replacement(
                add,
                sym_minus,
                my_args,
                fwd_only,
                [config.post_grad_custom_post_pass],
                extra_check=extra_check,
            )

            @torch.compile(dynamic=True)
            def foo(x, y):
                return x + y

            x = torch.rand([8, 1])
            y = torch.rand([10])

            self.assertEqual(foo(x, y), x + y)

            self.assertTrue(invoked)
            # we trace out the y.sym_size in replacement
            FileCheck().check("sym_size_int").check_same("num_users=2").check_same(
                "target=torch.ops.aten.sym_size"
            ).run(str(saved_graph))

    @inductor_config.patch(fx_graph_remote_cache=False)
    def test_match_with_mutation(self):
        counter = 0
        test_pass = PatternMatcherPass(pass_name="test")

        @register_graph_pattern(
            CallFunction(
                torch.add, KeywordArg("x"), CallFunction(torch.sin, KeywordArg("x"))
            ),
            pass_dict=test_pass,
        )
        def _test(match, x):
            nonlocal counter
            counter += 1

        def fn0(x, y):
            a = torch.sin(x)
            b = torch.add(x, a)
            return b

        def fn1(x, y):
            a = torch.sin(x)
            x.copy_(y)
            b = torch.add(x, a)
            return b

        def fn2(x, y):
            a = torch.sin(x)
            with torch.no_grad():
                b = torch.add(x, a)
            return b

        def fn3(x, y):
            a = torch.sin(x)
            with torch.autocast(GPU_TYPE):
                b = torch.add(x, a)
            return b

        def fn4(x, y):
            a = torch.sin(x)
            torch.manual_seed(1234)
            b = torch.add(x, a)
            return b

        def fn5(x, y):
            a = torch.sin(x)
            torch.add(y, 1, out=x)
            b = torch.add(x, a)
            return b

        args = [
            torch.randn(5, 5, device=GPU_TYPE),
            torch.randn(5, 5, device=GPU_TYPE),
        ]

        with unittest.mock.patch(
            "torch._inductor.fx_passes.pre_grad.config.pre_grad_fusion_options",
            {"test": {}},
        ), unittest.mock.patch(
            "torch._inductor.fx_passes.pre_grad.PRE_GRAD_FUSIONS",
            [],
        ), unittest.mock.patch(
            "torch._inductor.fx_passes.pre_grad.PRE_GRAD_PATTERNS",
            {"test": test_pass},
        ):
            for fn in (fn0, fn1, fn2, fn3, fn4, fn5):
                counter = 0
                expected = fn(*copy.deepcopy(args))
                actual = torch.compile(fn)(*copy.deepcopy(args))
                # should not match
                self.assertEqual(counter, int(fn is fn0))
                torch.testing.assert_close(actual, expected)

    def test_remove_pointless_clones(self):
        @torch.compile(fullgraph=True)
        def fn(a, b):
            return torch.mm(a, b).clone()

        result, (code) = run_and_get_code(fn, torch.randn(8, 8), torch.randn(8, 8))
        # clone would create a buf1
        self.assertIn("return (buf0, )", code[0])
        self.assertNotIn("async_compile.cpp", code[0])

    @expectedFailureXPU
    def test_unfuse_bias_addmm(self):
        args = [
            torch.randn(20, device=GPU_TYPE),
            torch.randn(10, 15, device=GPU_TYPE),
            torch.randn(15, 20, device=GPU_TYPE),
        ]

        @torch.compile()
        def fn(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b)

        _, (code) = run_and_get_code(fn, args[0], args[1], args[2])
        FileCheck().check("extern_kernels.addmm(").run(code[0])

        @torch.compile()
        def fn2(inp, a, b):
            return torch.nn.functional.gelu(torch.ops.aten.addmm(inp, a, b))

        _, (code) = run_and_get_code(fn2, args[0], args[1], args[2])
        FileCheck().check_not("extern_kernels.addmm(").run(code[0])

        @torch.compile()
        def fn2(inp, a, b):
            return torch.nn.functional.gelu(
                torch.ops.aten.addmm(inp, a, b).unsqueeze(0)
            )

        # hit the view path
        _, (code) = run_and_get_code(fn2, args[0], args[1], args[2])
        FileCheck().check_not("extern_kernels.addmm(").run(code[0])

    def test_serialized_patterns_up_to_date(self):
        import torch.utils._pytree as pytree
        from torch._inductor.fx_passes import joint_graph
        from torch._inductor.pattern_matcher import _known_precompiled_patterns

        # Ensure the patterns are loaded
        os.environ.pop("PYTORCH_GEN_PATTERNS", None)
        joint_graph.lazy_init()

        with torch._subclasses.FakeTensorMode() as mode:
            for (
                search_fn,
                example_inputs,
                trace_fn,
                scalar_workaround,
                search_fn_pattern,
            ) in _known_precompiled_patterns:
                # Because the example_inputs were saved as fake tensors in a
                # different FakeTensorMode we need to update them to our
                # FakeTensorMode().
                def remap_fake_tensor(x):
                    if isinstance(x, torch.Tensor):
                        return torch._subclasses.FakeTensor.from_tensor(x, mode)
                    return x

                example_inputs = pytree.tree_map(remap_fake_tensor, example_inputs)

                pattern = gen_pattern(
                    search_fn, example_inputs, trace_fn, scalar_workaround
                )
                pattern_pp = PatternPrettyPrinter.run(pattern)

                self.assertEqual(
                    pattern_pp,
                    PatternPrettyPrinter.run(search_fn_pattern),
                    msg=f"Found mismatched pattern {search_fn.__name__}. Run torchgen/fuse/gen_patterns.py",
                )

                # Since we've already checked that the serialized patterns match
                # lets verify the serializer by ensuring the generated patterns
                # also match (since search_fn_pattern is the serialized version
                # of search_fn).
                self.assertTrue(pattern.pattern_eq(search_fn_pattern))

    @skipIfXpu
    @xfailIfSM89
    @inductor_config.patch(
        {
            "triton.unique_kernel_names": "original_aten",
            "fx_graph_remote_cache": False,
            "max_autotune_gemm_backends": "TRITON",
        }
    )
    def test_original_aten_preserved_split_addmm(self):
        # addmm -> elementwise should be decomposed into mm -> add -> elementwise
        def fn(x, y, z):
            return torch.addmm(z, x, y).sin()

        args = [
            torch.randn(16, 24, device=GPU_TYPE),
            torch.randn(24, 32, device=GPU_TYPE),
            torch.randn(16, 32, device=GPU_TYPE),
        ]

        counters.clear()

        opt_fn = torch.compile(fn, mode="max-autotune")
        ret, code = run_and_get_code(opt_fn, *args)
        self.assertEqual(counters["inductor"]["pattern_matcher_count"], 1)

        # The mm kernel should use a template (because we set max_autotune_gemm_backends = TRITON).
        # Its name should contain `addmm` because `addmm` was the original aten op where the mm came from.
        FileCheck().check_not("extern_kernels.addmm(").check(
            "def triton_tem_fused_addmm"
        ).run(code[0])

    @inductor_config.patch(fx_graph_remote_cache=False)
    def test_match_equivalent_function_invocations1(self):
        counter = 0
        test_pass = PatternMatcherPass()

        args = [
            torch.randn(20, device=GPU_TYPE),
            torch.randn(10, 15, device=GPU_TYPE),
            torch.randn(15, 20, device=GPU_TYPE),
        ]

        def f0(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b)

        def f1(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b, beta=1.0)

        def f2(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b, beta=1.0, alpha=1.0)

        # This graph pattern should successfully match all of the above functions
        @register_graph_pattern(
            CallFunction(
                torch.ops.aten.addmm,
                Arg(),
                Arg(),
                Arg(),
                beta=KeywordArg("beta"),
                alpha=KeywordArg("alpha"),
            ),
            pass_dict=test_pass,
        )
        def addmm_replacement(match: Match, inp, mat1, mat2, beta, alpha):
            nonlocal counter
            counter += 1

            def repl(inp, x1, x2):
                return (x1 @ x2) * alpha + inp * beta

            with V.fake_mode:
                match.replace_by_example(repl, [inp, mat1, mat2])

        with unittest.mock.patch(
            "torch._inductor.fx_passes.post_grad.pass_patterns",
            torch._inductor.fx_passes.post_grad.pass_patterns + [test_pass],
        ):
            for fn in (f0, f1, f2):
                counter = 0
                expected = fn(*copy.deepcopy(args))
                opt_fn = torch.compile(fn)
                actual, (code) = run_and_get_code(opt_fn, args[0], args[1], args[2])
                # pattern should match
                self.assertEqual(counter, 1)
                torch.testing.assert_close(actual, expected)
                # addmm should be replaced
                FileCheck().check_not("extern_kernels.addmm(").run(code[0])

    @inductor_config.patch(fx_graph_remote_cache=False)
    def test_match_equivalent_function_invocations2(self):
        counter = 0
        test_pass = PatternMatcherPass()

        args = [
            torch.randn(20, device=GPU_TYPE),
            torch.randn(10, 15, device=GPU_TYPE),
            torch.randn(15, 20, device=GPU_TYPE),
        ]

        def f0(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b)

        def f1(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b, beta=1.0)

        def f2(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b, beta=1.0, alpha=1.0)

        # This graph pattern should only match f0
        @register_graph_pattern(
            CallFunction(torch.ops.aten.addmm, Arg(), Arg(), Arg()),
            pass_dict=test_pass,
        )
        def addmm_replacement(match: Match, inp, mat1, mat2):
            nonlocal counter
            counter += 1

            def repl(inp, x1, x2):
                return x1 @ x2 + inp

            with V.fake_mode:
                match.replace_by_example(repl, [inp, mat1, mat2])

        with unittest.mock.patch(
            "torch._inductor.fx_passes.post_grad.pass_patterns",
            torch._inductor.fx_passes.post_grad.pass_patterns + [test_pass],
        ):
            for fn in (f0, f1, f2):
                counter = 0
                expected = fn(*copy.deepcopy(args))
                actual = torch.compile(fn)(*copy.deepcopy(args))
                self.assertEqual(counter, 1)
                torch.testing.assert_close(actual, expected)

    @inductor_config.patch(fx_graph_remote_cache=False)
    def test_match_equivalent_function_invocations3(self):
        counter = 0
        test_pass = PatternMatcherPass()

        args = [
            torch.randn(20, device=GPU_TYPE),
            torch.randn(10, 15, device=GPU_TYPE),
            torch.randn(15, 20, device=GPU_TYPE),
        ]

        def f0(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b)

        def f1(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b, beta=1.0)

        def f2(inp, a, b):
            return torch.ops.aten.addmm(inp, a, b, beta=1.0, alpha=1.0)

        # This graph pattern should only match f1
        @register_graph_pattern(
            CallFunction(
                torch.ops.aten.addmm, Arg(), Arg(), Arg(), beta=KeywordArg("beta")
            ),
            pass_dict=test_pass,
        )
        def addmm_replacement(match: Match, inp, mat1, mat2, beta):
            nonlocal counter
            counter += 1

            def repl(inp, x1, x2):
                return x1 @ x2 + inp

            with V.fake_mode:
                match.replace_by_example(repl, [inp, mat1, mat2])

        with unittest.mock.patch(
            "torch._inductor.fx_passes.post_grad.pass_patterns",
            torch._inductor.fx_passes.post_grad.pass_patterns + [test_pass],
        ):
            for fn in (f0, f1, f2):
                counter = 0
                expected = fn(*copy.deepcopy(args))
                actual = torch.compile(fn)(*copy.deepcopy(args))
                self.assertEqual(counter, 1)
                torch.testing.assert_close(actual, expected)

    def test_stable_topological_sort(self):
        def fn1(a, b):
            return a + b

        graph = torch.fx.Graph()
        a = graph.placeholder("x")
        b = graph.placeholder("y")
        c = graph.call_function(fn1, (a, b))
        stable_topological_sort(graph)
        self.assertEqual(list(graph.nodes), [a, b, c])

        graph = torch.fx.Graph()
        b = graph.placeholder("y")
        a = graph.placeholder("x")
        c = graph.call_function(fn1, (a, b))
        stable_topological_sort(graph)
        self.assertEqual(list(graph.nodes), [b, a, c])

        graph = torch.fx.Graph()
        a = graph.placeholder("x")
        b = graph.placeholder("y")
        c = graph.call_function(fn1, (b, a))
        c.append(a)
        stable_topological_sort(graph)
        self.assertEqual(list(graph.nodes), [b, a, c])

    def test_scaled_softmax(self):
        def mul_softmax(a, b):
            return F.softmax(a * b, dim=0)

        def div_softmax(x, inv_scale):
            return F.softmax(x / inv_scale, dim=0)

        x = torch.randn(10, 10)
        scale = 1e6
        inv_scale = 1 / scale
        self.common(mul_softmax, (x, scale), 1, 3)
        self.common(mul_softmax, (scale, x), 1, 3)
        self.common(div_softmax, (x, inv_scale), 1, 3)

        scale = torch.randn(10) * 1e6
        inv_scale = 1 / scale
        self.common(mul_softmax, (x, scale), 1, 3)
        self.common(mul_softmax, (scale, x), 1, 3)
        self.common(div_softmax, (x, inv_scale), 1, 3)

        scale = torch.randn(1, 10) * 1e6
        inv_scale = 1 / scale
        self.common(mul_softmax, (x, scale), 1, 3)
        self.common(mul_softmax, (scale, x), 1, 3)
        self.common(div_softmax, (x, inv_scale), 1, 3)

        # Test matching with type promotion
        x = torch.randn(10, 10, dtype=torch.bfloat16)
        scale = torch.randn(10, dtype=torch.bfloat16) * 1e6
        inv_scale = 1 / scale
        self.common(mul_softmax, (x, scale), 1, 4, reference_in_float=True)
        self.common(mul_softmax, (scale, x), 1, 4, reference_in_float=True)
        self.common(div_softmax, (x, inv_scale), 1, 4, reference_in_float=True)

        # No match if scale changes in softmax dim
        scale = torch.randn(10, 10)
        self.common(mul_softmax, (x, scale), 0, 0)
        self.common(mul_softmax, (scale, x), 0, 0)
        self.common(div_softmax, (x, scale), 0, 0)

    def test_mutation_op_matching(self):
        def check(type, func_name, args, kwargs, expect=True):
            assert type in ["call_function", "call_method"]
            graph = torch.fx.Graph()
            getattr(graph, type)(func_name, args, kwargs)
            res = is_mutation_op(next(iter(graph.nodes)))
            if expect:
                self.assertTrue(res)
            else:
                self.assertFalse(res)

        t = torch.randn(1)
        check("call_function", torch._C._set_grad_enabled, (False,), {})
        check("call_method", "copy_", (t, t), {})
        check("call_method", "relu_", (t,), {})
        check("call_function", torch.manual_seed, (0,), {})
        check("call_function", torch.ops.aten.set_.source_Tensor, (t, t), {})
        check(
            "call_function",
            torch.amp.autocast_mode._enter_autocast,
            (GPU_TYPE, None, True, None),
            {},
        )
        check("call_function", torch.amp.autocast_mode._exit_autocast, (None,), {})
        check(
            "call_function",
            torch.ops._c10d_functional.all_gather_into_tensor_out,
            (t, 2, "0"),
            {"out": t},
        )
        check("call_function", torch.ops.inductor.resize_storage_bytes_, (t, 0), {})
        check(
            "call_function",
            torch.ops.inductor.resize_storage_bytes_.default,
            (t, 0),
            {},
        )
        check(
            "call_function",
            torch.ops.fsdp.split_with_sizes_copy,
            (t, [64, 128, 8, 8]),
            {"dim": 1, "out": [t, t, t, t]},
        )
        check("call_function", torch.ops.fsdp.copy_, (t, t), {})
        check(
            "call_function", torch.ops.aten.__rshift__.Scalar, (t, 2), {}, expect=False
        )
        check(
            "call_function",
            torch.ops._c10d_functional.all_gather_into_tensor,
            (t, 2, "0"),
            {},
            expect=False,
        )

        @torch.library.custom_op("vllm::fused_rms_norm_quant_static", mutates_args=[])
        def fused_rms_norm_quant_static(out: torch.Tensor, input: torch.Tensor) -> None:
            pass

        check(
            "call_function",
            torch.ops.vllm.fused_rms_norm_quant_static,
            (t, t),
            {},
            expect=False,
        )

    def test_multioutput_register_replacement(self):
        @torch.library.custom_op(
            "vllm::fused_rms_norm_quant_static", mutates_args=["result", "scale"]
        )
        def fused_rms_norm_quant_static(
            result: torch.Tensor,
            input: torch.Tensor,
            weight: torch.Tensor,
            scale: torch.Tensor,
            azp: torch.Tensor,
            epsilon: float,
        ) -> None:
            print("vllm::fused_rms_norm_quant_static")
            result_rms = torch.mul(input, weight) + epsilon
            result = torch.mul(result_rms, scale).to(torch.int8)
            scale.fill_(0.5)

        @torch.library.custom_op("vllm::rms_norm", mutates_args=["result"])
        def rms_norm(
            result: torch.Tensor,
            input: torch.Tensor,
            weight: torch.Tensor,
            epsilon: float,
        ) -> None:
            # bogus implementation doesn't matter
            result = torch.mul(input, weight) + epsilon

        @torch.library.custom_op(
            "vllm::static_scaled_int8_quant", mutates_args=["result", "scale"]
        )
        def static_scaled_int8_quant(
            result: torch.Tensor,
            input: torch.Tensor,
            scale: torch.Tensor,
            azp: Optional[torch.Tensor] = None,
        ) -> None:
            # bogus implementation doesn't matter
            result = torch.mul(input, scale).to(torch.int8)
            scale.fill_(0.5)

        def rms_pattern_static(
            result: torch.Tensor,
            result_rms: torch.Tensor,
            input: torch.Tensor,
            weight: torch.Tensor,
            scale: torch.Tensor,
        ):
            at1 = auto_functionalized(
                torch.ops.vllm.rms_norm.default,
                result=result_rms,
                input=input,
                weight=weight,
                epsilon=1e-6,
            )
            at2 = auto_functionalized(
                torch.ops.vllm.static_scaled_int8_quant.default,
                result=result,
                input=at1[1],
                scale=scale,
                azp=None,
            )

            return at2[1], at2[2]

        def rms_replacement_static(
            result: torch.Tensor,
            result_rms: torch.Tensor,
            input: torch.Tensor,
            weight: torch.Tensor,
            scale: torch.Tensor,
        ):
            at = auto_functionalized(
                torch.ops.vllm.fused_rms_norm_quant_static.default,
                result=result,
                input=input,
                weight=weight,
                epsilon=1e-6,
                scale=scale,
                azp=None,
            )
            return at[1], at[2]

        def empty_bf16(*args, **kwargs):
            return torch.empty(*args, **kwargs, dtype=torch.bfloat16)

        def empty_int8(*args, **kwargs):
            return torch.empty(*args, **kwargs, dtype=torch.int8)

        my_patterns = PatternMatcherPass()
        inputs = [
            empty_int8(5, 4),
            empty_bf16(5, 4),
            empty_bf16(5, 4),
            empty_bf16(5, 1),
            torch.empty(1, 1),
        ]
        register_replacement(
            rms_pattern_static, rms_replacement_static, inputs, fwd_only, my_patterns
        )

        def custom_pass(graph: torch.fx.Graph) -> torch.fx.Graph:
            count = my_patterns.apply(graph)
            # print(f"Count: {count}")
            graph.eliminate_dead_code()
            # graph.print_tabular()
            return graph

        def custom_backend(
            graph: torch.fx.GraphModule, example_inputs: List[torch.Tensor]
        ) -> Callable:
            from torch._inductor import config

            current_config = config.shallow_copy_dict()
            from torch._inductor.compile_fx import compile_fx

            current_config["post_grad_custom_post_pass"] = custom_pass
            return compile_fx(graph, example_inputs, config_patches=current_config)

        @torch.compile(backend=custom_backend)
        def my_func_static(x, w, epsilon):
            quant_result = torch.empty_like(x, dtype=torch.int8)
            result_rms = torch.empty_like(x, dtype=torch.bfloat16)
            scale = torch.ones((1, 1))

            x = x.to(torch.bfloat16)
            w = w.to(torch.bfloat16)

            quant_result, scale = rms_pattern_static(
                result=quant_result,
                result_rms=result_rms,
                input=x,
                weight=w,
                scale=scale,
            )

            return quant_result, scale

        inputs = [torch.empty((5, 4)), torch.empty((5, 1)), 1e-6]
        # print(my_func_static(*inputs))
        test, (code,) = run_and_get_code(my_func_static, *inputs)
        self.assertTrue("static_scaled_int8_quant" not in code)


if __name__ == "__main__":
    if IS_LINUX and HAS_GPU:
        run_tests()