# Owner(s): ["oncall: export"]
import copy
import re
import tempfile
import unittest

import torch
from torch._subclasses.fake_tensor import FakeTensorMode
from torch.export import Dim, draft_export, export
from torch.export._draft_export import FailureType
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.testing import FileCheck
from torch.testing._internal.common_utils import IS_WINDOWS, run_tests, TestCase
from torch.testing._internal.torchbind_impls import (
    _empty_tensor_queue,
    init_torchbind_implementations,
)
from torch.utils._pytree import tree_leaves


class TestDraftExport(TestCase):
    def setUp(self):
        super().setUp()
        init_torchbind_implementations()

        self.torch_bind_ops = [
            torch.ops._TorchScriptTesting.queue_pop,
            torch.ops._TorchScriptTesting.queue_push,
            torch.ops._TorchScriptTesting.queue_size,
        ]

    def tearDown(self):
        return

    def test_missing_meta_kernel_custom_op_basic(self):
        with torch.library._scoped_library("mylib", "FRAGMENT"):

            @torch.library.custom_op("mylib::foo2", mutates_args={})
            def foo2_impl(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
                return a + b

            class M(torch.nn.Module):
                def forward(self, a, b):
                    res = torch.ops.mylib.foo2(a, b)
                    return res

            inp = (torch.ones(3, 3), torch.ones(3, 3))

            ep = draft_export(M(), inp)
            report = ep._report

            self.assertEqual(len(report.failures), 1)
            self.assertEqual(
                report.failures[0].failure_type, FailureType.MISSING_FAKE_KERNEL
            )

            inp = (torch.randn(3, 3), torch.randn(3, 3))
            self.assertEqual(ep.module()(*inp), M()(*inp))

            with torch._library.fake_profile.unsafe_generate_fake_kernels(
                report.op_profiles
            ):
                ep.run_decompositions()

    def test_missing_meta_kernel_impl(self):
        with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
            torch.library.define(
                "mylib::foo",
                "(Tensor a, Tensor b) -> Tensor",
                tags=torch.Tag.pt2_compliant_tag,
                lib=lib,
            )

            @torch.library.impl("mylib::foo", "cpu", lib=lib)
            def foo_impl(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
                return a + b

            class M(torch.nn.Module):
                def forward(self, a, b):
                    res = torch.ops.mylib.foo(a, b)
                    res = torch.ops.mylib.foo(res, b)
                    return res

            inp = (torch.ones(3, 3), torch.ones(3, 3))

            ep = draft_export(M(), inp)
            report = ep._report

            self.assertEqual(len(report.failures), 1)
            self.assertEqual(
                report.failures[0].failure_type, FailureType.MISSING_FAKE_KERNEL
            )

            inp = (torch.randn(3, 3), torch.randn(3, 3))
            self.assertEqual(ep.module()(*inp), M()(*inp))

            self.assertEqual(len(report.op_profiles), 1)
            self.assertEqual(len(report.op_profiles["mylib.foo.default"]), 1)
            print(report.op_profiles)

            with torch._library.fake_profile.unsafe_generate_fake_kernels(
                report.op_profiles
            ):
                ep = ep.run_decompositions()
            self.assertEqual(ep.module()(*inp), M()(*inp))

    def test_missing_meta_kernel_custom_op_multiple_profiles(self):
        with torch.library._scoped_library("mylib", "FRAGMENT"):

            @torch.library.custom_op("mylib::foo3", mutates_args={})
            def foo3_impl(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
                return a + b

            class M(torch.nn.Module):
                def forward(self, a, b, c, d):
                    res1 = torch.ops.mylib.foo3(a, b)
                    res2 = torch.ops.mylib.foo3(c, d)
                    return res1, res2

            inp = (
                torch.ones(3, 4),
                torch.ones(3, 4),
                torch.ones(2, 3, 4),
                torch.ones(2, 3, 4),
            )

            ep = draft_export(M(), inp)
            report = ep._report

            self.assertEqual(len(report.failures), 1)
            self.assertEqual(
                report.failures[0].failure_type, FailureType.MISSING_FAKE_KERNEL
            )
            self.assertEqual(len(report.op_profiles), 1)
            self.assertEqual(len(report.op_profiles["mylib.foo3.default"]), 2)

            with torch._library.fake_profile.unsafe_generate_fake_kernels(
                report.op_profiles
            ):
                ep.run_decompositions()

    def test_missing_meta_kernel_custom_op_update_profile(self):
        with torch.library._scoped_library("mylib", "FRAGMENT"):

            @torch.library.custom_op("mylib::foo8", mutates_args={})
            def foo8_impl(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
                return a + b

            class M(torch.nn.Module):
                def forward(self, a, b):
                    res = torch.ops.mylib.foo8(a, b)
                    return res

            inp = (
                torch.ones(3, 4),
                torch.ones(3, 4),
            )

            ep = draft_export(M(), inp)
            report = ep._report
            self.assertEqual(len(report.op_profiles), 1)
            self.assertEqual(len(report.op_profiles["mylib.foo8.default"]), 1)

            new_inp = (
                torch.ones(2, 3, 4),
                torch.ones(2, 3, 4),
            )

            with torch._library.fake_profile.unsafe_generate_fake_kernels(
                report.op_profiles
            ):
                with FakeTensorMode(allow_non_fake_inputs=True, shape_env=ShapeEnv()):
                    torch.ops.mylib.foo8(*inp)
                    with self.assertRaisesRegex(
                        RuntimeError, "no profiles match the given inputs"
                    ):
                        torch.ops.mylib.foo8(*new_inp)

                ep = draft_export(M(), new_inp)

            report = ep._report
            self.assertEqual(len(report.op_profiles), 1)
            self.assertEqual(len(report.op_profiles["mylib.foo8.default"]), 1)

            with (
                torch._library.fake_profile.unsafe_generate_fake_kernels(
                    report.op_profiles
                ),
                FakeTensorMode(allow_non_fake_inputs=True, shape_env=ShapeEnv()),
            ):
                torch.ops.mylib.foo8(*new_inp)

                # Existing registration has been updated to match the new
                # profile traced with draft-export
                with self.assertRaisesRegex(
                    RuntimeError, "no profiles match the given inputs"
                ):
                    torch.ops.mylib.foo8(*inp)

    @unittest.skipIf(not torch.cuda.is_available(), "Requires cuda")
    def test_missing_meta_kernel_guard(self):
        with torch.library._scoped_library("mylib", "FRAGMENT"):

            @torch.library.custom_op("mylib::foo4", mutates_args={})
            def foo4_impl(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
                return a + b

            class M(torch.nn.Module):
                def forward(self, a, b):
                    res1 = torch.ops.mylib.foo4(a, b)
                    return res1

            inp = (
                torch.ones(3, 4),
                torch.ones(3, 4),
            )

            ep = draft_export(
                M(),
                inp,
                dynamic_shapes={
                    "a": {0: Dim.DYNAMIC, 1: Dim.DYNAMIC},
                    "b": {0: Dim.DYNAMIC, 1: Dim.DYNAMIC},
                },
            )

            inp = (torch.randn(2, 3), torch.randn(2, 3))
            self.assertEqual(ep.module()(*inp), M()(*inp))
            m = ep.module()
            with self.assertRaisesRegex(RuntimeError, "Tensor device mismatch!"):
                bad_device_inps = (
                    torch.randn(2, 3, device=torch.device("cuda")),
                    torch.randn(2, 3, device=torch.device("cuda")),
                )
                m(*bad_device_inps)

            with self.assertRaisesRegex(RuntimeError, "Tensor dtype mismatch!"):
                bad_dtype_inps = (
                    torch.randn(2, 3, dtype=torch.float16),
                    torch.randn(2, 3, dtype=torch.float16),
                )
                m(*bad_dtype_inps)

    def test_fake_infer_dense_in_memory_check(self):
        with torch.library._scoped_library("mylib", "FRAGMENT"):

            @torch.library.custom_op("mylib::foo5", mutates_args={})
            def foo5_impl(a: torch.Tensor) -> torch.Tensor:
                return a * 2

            @torch.library.custom_op("mylib::foo6", mutates_args={})
            def foo6_impl(a: torch.Tensor) -> torch.Tensor:
                return (a * 2)[:, :-1, :-1]  # not dense in memory

            @torch.library.custom_op("mylib::foo7", mutates_args={})
            def foo7_impl(a: torch.Tensor) -> torch.Tensor:
                return (a * 2)[:, 1:-1, :]  # non-zero storage offset

            class Foo(torch.nn.Module):
                def forward(self, x, opt):
                    if opt == 0:
                        return torch.ops.mylib.foo5(x)
                    elif opt == 1:
                        return torch.ops.mylib.foo6(x)
                    else:
                        return torch.ops.mylib.foo7(x)

            draft_export(Foo(), (torch.randn(80, 4, 4), 0))
            draft_export(Foo(), (torch.randn(80, 1, 4), 0))
            draft_export(Foo(), (torch.randn(1, 4, 1, 1, 4, 1, 4), 0))
            with self.assertRaisesRegex(
                RuntimeError,
                "a return was not dense in memory",
            ):
                draft_export(Foo(), (torch.randn(4, 6, 8), 1))
            with self.assertRaisesRegex(
                RuntimeError,
                "a return has a non-zero storage offset",
            ):
                draft_export(Foo(), (torch.randn(4, 6, 8), 2))

    def test_data_dependent_failure(self):
        with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
            torch.library.define(
                "mylib::foo1",
                "(Tensor a, Tensor b) -> Tensor",
                tags=torch.Tag.pt2_compliant_tag,
                lib=lib,
            )

            @torch.library.impl("mylib::foo1", "cpu", lib=lib)
            def foo_impl(a, b):
                return a + b

            class M(torch.nn.Module):
                def forward(self, a, b, c):
                    res = torch.ops.mylib.foo1(a, b)

                    c_item = c.item()
                    return res[:c_item]

            inp = (torch.ones(3, 3), torch.ones(3, 3), torch.tensor(3))

            ep = draft_export(M(), inp)
            report = ep._report
            self.assertTrue(len(report.failures) > 0)
            self.assertEqual(
                report.failures[0].failure_type, FailureType.MISSING_FAKE_KERNEL
            )
            self.assertEqual(
                report.failures[1].failure_type, FailureType.DATA_DEPENDENT_ERROR
            )

            inp = (torch.randn(3, 3), torch.randn(3, 3), torch.tensor(2))
            self.assertEqual(ep.module()(*inp), M()(*inp))

    def test_unbacked_div_mod_replacement(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = torch.zeros(x.item())
                x = x.unsqueeze(0).repeat(10, 2)
                return x.view(-1, 2, 2345)

        ep = draft_export(M(), (torch.tensor([938]),))
        report = ep._report
        self.assertEqual(len(report.failures), 0)

    def test_dedup_data_dependent_failure(self):
        class M(torch.nn.Module):
            def forward(self, x, y, z):
                res = 0
                for v in [x, y]:
                    b = v.item()
                    if b > 10:
                        res += v * b
                    else:
                        res += v + b

                return z * res

        inp = (torch.tensor(5), torch.tensor(3), torch.tensor(2))

        ep = draft_export(M(), inp)
        report = ep._report
        self.assertEqual(len(report.failures), 1)
        self.assertEqual(
            report.failures[0].failure_type, FailureType.DATA_DEPENDENT_ERROR
        )

        inp = (torch.tensor(4), torch.tensor(2), torch.tensor(6))
        self.assertEqual(ep.module()(*inp), M()(*inp))

        # the fake tensors on node.meta["val"] should have real_tensor
        gm = ep.module()
        tensors = [
            node.meta.get("val").real_tensor
            for node in gm.graph.nodes
            if node.op == "placeholder"
        ]
        self.assertTrue(all(isinstance(t, torch.Tensor) for t in tensors))

    def test_complex_data_dependent_expr(self):
        class M(torch.nn.Module):
            def forward(self, x, y):
                a = x.item()
                a = -a
                a = a // 3
                a = a + 5

                z = torch.cat([y, y])

                return z[:a]

        ep = draft_export(
            M(),
            (torch.tensor(6), torch.randn(5)),
            dynamic_shapes={"x": None, "y": {0: Dim.DYNAMIC}},
        )
        report = ep._report
        self.assertTrue(len(report.failures) > 0)
        self.assertEqual(
            report.failures[0].failure_type, FailureType.DATA_DEPENDENT_ERROR
        )
        for _ep in [ep, ep.run_decompositions()]:
            # unbacked bindings
            unbacked_binding_symbols = set()
            for node in _ep.graph.nodes:
                if bindings := node.meta.get("unbacked_bindings"):
                    unbacked_binding_symbols.update(bindings.keys())
            self.assertEqual(len(unbacked_binding_symbols), 1)

    def test_offsets(self):
        class M(torch.nn.Module):
            def forward(self, x):
                a = x.item()
                if a == 0:
                    raise RuntimeError("bad")
                return x * a

        inp = (torch.tensor(3),)
        draft_export(M(), inp)

    def test_shape_failure(self):
        class M(torch.nn.Module):
            def forward(self, a):
                assert a.shape[0] == 3
                return a * a

        inp = (torch.ones(3, 3),)

        ep = draft_export(
            M(),
            inp,
            dynamic_shapes={"a": {0: Dim("a0")}},
            prefer_deferred_runtime_asserts_over_guards=True,
        )
        report = ep._report

        self.assertEqual(len(report.failures), 1)
        self.assertEqual(report.failures[0].failure_type, FailureType.GUARD_ADDED)

        inp = (torch.randn(3, 3),)
        self.assertEqual(ep.module()(*inp), M()(*inp))

        inp = (torch.randn(4, 3),)
        with self.assertRaisesRegex(
            AssertionError,
            re.escape("Guard failed: a.size()[0] <= 3"),
        ):
            # expected <= 3, but got 4
            ep.module()(*inp)

    def test_side_effect1(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.register_buffer("a", torch.tensor(2))

            def forward(self, b):
                a_item = self.a.item()
                if a_item == 2:
                    res = a_item * b
                else:
                    res = (a_item + 1) * b

                self.a.add_(1)
                a_item = self.a.item()

                if a_item == 3:
                    res = a_item * res
                else:
                    res = (a_item + 1) * res
                return res

        inp = (torch.ones(3, 3),)
        mod = M()
        ep = draft_export(mod, inp)
        self.assertEqual(mod.a, torch.tensor(2))
        FileCheck().check_count("torch.ops.aten.add.default", 0, exactly=True).run(
            ep.graph_module.code
        )

    def test_side_effect_inps(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()

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

        inp = (torch.ones(3, 3),)
        ep = draft_export(M(), inp)
        report = ep._report
        self.assertTrue(report.successful())
        self.assertEqual(inp[0], torch.ones(3, 3))

    def test_masked_linear(self):
        class M(torch.nn.Module):
            def forward(self, x, mask, weight, bias):
                masked = x[mask != 0, :, :]
                return torch.nn.functional.linear(masked, weight, bias)

        x = torch.zeros(10)
        x[0] += 1
        inp = (torch.randn(10, 8, 7), x, torch.randn(25, 7), torch.randn(25))
        draft_ep = draft_export(M(), inp)
        ep = export(M(), inp)
        self.assertEqual(draft_ep.module()(*inp), ep.module()(*inp))
        x[2] += 1
        x[3] += 1
        self.assertEqual(draft_ep.module()(*inp), ep.module()(*inp))

    def test_torchbind(self):
        class Model(torch.nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.linear = torch.nn.Linear(2, 2)

            def forward(self, tq, x):
                x_cos = tq.pop() + tq.float_size() + self.linear(x)
                if tq.is_empty():
                    x_sin = self.linear(tq.pop()) - tq.size() + x
                else:
                    x_sin = tq.pop() + tq.size() + x
                return x_sin, x_cos, tq

        mod = Model()
        tq = _empty_tensor_queue()
        tq2 = copy.deepcopy(tq)
        a = torch.randn(2, 2)
        b = torch.randn(2, 2)
        tq.push(a)
        tq.push(b)
        tq3 = copy.deepcopy(tq)
        inp = (tq, torch.randn(2, 2))
        ep = draft_export(mod, inp)
        report = ep._report
        self.assertTrue(report.successful())
        self.assertEqual(tq2.size(), 0)
        self.assertEqual(tq3.size(), 2)
        self.assertEqual(tq.size(), 2)

    def test_override_size_and_dtype_mismatched_fake_kernels(self):
        with torch.library._scoped_library("mylib", "FRAGMENT"):

            class M(torch.nn.Module):
                def forward(self, a):
                    return torch.ops.mylib.foo9(a)

            @torch.library.custom_op("mylib::foo9", mutates_args={})
            def foo(a: torch.Tensor) -> list[torch.Tensor]:
                x = a * 2
                y = a.repeat(2, 2)
                z = a.to(torch.bfloat16)
                return [x, y, z]

            @torch.library.register_fake("mylib::foo9")
            def foo_fake_impl(a):
                x = torch.empty_like(a)  # good
                y = torch.empty_like(a)  # size mismatch
                z = torch.empty_like(a)  # dtype mismatch
                return [x, y, z]

            mod = M()
            inputs = (torch.randn(3, 3),)
            with self.assertRaises(RuntimeError):
                with torch._functorch.config.patch(
                    fake_tensor_propagate_real_tensors=True
                ):
                    export(mod, inputs, strict=True)

            ep = draft_export(mod, inputs)
            report = ep._report
            for ep_out, eager_out in zip(ep.module()(*inputs), mod(*inputs)):
                self.assertTrue(torch.allclose(ep_out, eager_out))
                self.assertEqual(ep_out.dtype, eager_out.dtype)

            self.assertEqual(len(report.failures), 2)
            self.assertEqual(
                report.failures[0].failure_type, FailureType.MISMATCHED_FAKE_KERNEL
            )
            self.assertEqual(
                report.failures[1].failure_type, FailureType.MISMATCHED_FAKE_KERNEL
            )
            self.assertEqual(
                sorted([f.data["reason"] for f in report.failures]),
                [
                    "Dtypes torch.bfloat16 and torch.float32 are not equal!",
                    "mismatch between fake value 3 and real value 6 ",
                ],
            )

            with torch._library.fake_profile.unsafe_generate_fake_kernels(
                report.op_profiles
            ):
                ep.run_decompositions()

    def test_override_incorrectly_aliasing_kernel(self):
        with torch.library._scoped_library("mylib", "FRAGMENT"):

            @torch.library.custom_op("mylib::foo10", mutates_args={})
            def foo(a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
                return a * 2, a + 2

            @torch.library.register_fake("mylib::foo10")
            def foo_fake_impl(a):
                return a, torch.empty_like(a)  # incorrectly aliasing

            class M(torch.nn.Module):
                def forward(self, a):
                    return torch.ops.mylib.foo10(a)

            mod = M()
            inputs = (torch.randn(3, 3),)
            with self.assertRaisesRegex(
                RuntimeError,
                "Real tensor propagation found an aliasing mismatch",
            ):
                with torch._functorch.config.patch(
                    fake_tensor_propagate_real_tensors=True
                ):
                    export(mod, inputs, strict=True)

            ep = draft_export(mod, inputs)
            report = ep._report
            for ep_out, eager_out in zip(
                tree_leaves(ep.module()(*inputs)), tree_leaves(mod(*inputs))
            ):
                self.assertTrue(torch.allclose(ep_out, eager_out))
                self.assertEqual(ep_out.dtype, eager_out.dtype)

            self.assertEqual(len(report.failures), 1)
            self.assertEqual(
                report.failures[0].failure_type, FailureType.MISMATCHED_FAKE_KERNEL
            )
            self.assertTrue(
                "Mismatched aliasing spec between fake kernel and real kernel"
                in report.failures[0].data["reason"]
            )

    def test_override_mismatched_fake_kernel_with_unbacked_symbols(self):
        with torch.library._scoped_library("mylib", "FRAGMENT"):

            @torch.library.custom_op("mylib::foo11", mutates_args={})
            def foo11(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
                return a[b.item()].to(torch.bfloat16)

            @torch.library.register_fake("mylib::foo11")
            def foo_fake_impl(a, b):
                ctx = torch.library.get_ctx()
                u = ctx.new_dynamic_size()
                return torch.empty(u, a.shape[1], dtype=a.dtype)

            class M(torch.nn.Module):
                def forward(self, a, b):
                    return torch.ops.mylib.foo11(a, b)

            mod = M()
            inputs = (torch.randn(100, 4), torch.tensor(10))

            ep = draft_export(mod, inputs)

            report = ep._report
            for ep_out, eager_out in zip(ep.module()(*inputs), mod(*inputs)):
                self.assertTrue(torch.allclose(ep_out, eager_out))
                self.assertEqual(ep_out.dtype, eager_out.dtype)

            self.assertEqual(len(report.failures), 1)
            self.assertEqual(
                report.failures[0].failure_type, FailureType.MISMATCHED_FAKE_KERNEL
            )
            self.assertEqual(
                report.failures[0].data["reason"],
                "Dtypes torch.bfloat16 and torch.float32 are not equal!",
            )
            with torch._library.fake_profile.unsafe_generate_fake_kernels(
                report.op_profiles
            ):
                ep.run_decompositions()

    # https://github.com/pytorch/pytorch/issues/140625
    @unittest.skipIf(IS_WINDOWS, "aoti_compile_and_package not supported on Windows")
    def test_constantify_unbacked_symbol(self):
        class M(torch.nn.Module):
            def forward(self, x, y):
                xt = torch.tensor(x.shape)
                u0 = xt[0].item()
                return y * torch.arange(u0)

        mod = M()
        example_inputs = (torch.randn(3, 5), torch.randn(3))
        draft_ep = draft_export(mod, example_inputs)
        with tempfile.NamedTemporaryFile(suffix=".pt2") as f:
            torch._inductor.aoti_compile_and_package(
                draft_ep,
                package_path=f.name,
            )

    @unittest.skipIf(
        not torch.cuda.is_available()
        or torch.cuda.get_device_properties(0).total_memory < 2**28,
        "Requires 16 MB GPU memory to pass the test; setting it higher to catch violations",
    )
    def test_cuda_memory_usage(self):
        # This used to OOM
        class Foo(torch.nn.Module):
            def forward(self, x):
                for _ in range(100):
                    x = x + 1e-3
                return x

        # measure base usage
        device = torch.device("cuda:0")
        torch.cuda.reset_peak_memory_stats()
        base_usage = torch.cuda.memory_allocated(device)

        # usage with input tensor allocated
        x = torch.randn(2**10, 2**10).to(device)
        x_usage = torch.cuda.memory_allocated(device)

        # draft export peak memory usage
        draft_export(Foo(), (x,), strict=False)
        peak_mem_usage = torch.cuda.memory_stats(device)["allocated_bytes.all.peak"]

        # right now it's actually exactly 4x;
        # I guess original tensor, 2 tensors per add op, 1 for clone stored in node.meta["val"]
        self.assertTrue((peak_mem_usage - base_usage) <= (x_usage - base_usage) * 4.0)


if __name__ == "__main__":
    run_tests()