# Owner(s): ["oncall: jit"]

import torch
from torch._C import parse_ir
from torch.testing._internal.common_utils import TemporaryFileName
from torch.testing._internal.jit_utils import JitTestCase


if __name__ == "__main__":
    raise RuntimeError(
        "This test file is not meant to be run directly, use:\n\n"
        "\tpython test/test_jit.py TESTNAME\n\n"
        "instead."
    )


class TestAliasAnalysis(JitTestCase):
    def test_becomes_wildcard_annotations(self):
        graph_str = """
        graph(%a.1 : Tensor, %b.1 : Tensor):
            %11 : NoneType = prim::Constant()
            %8 : int = prim::Constant[value=0]()
            %7 : int = prim::Constant[value=1]()
            %x.1 : Tensor = aten::add(%a.1, %b.1, %7)
            %y.1 : Tensor[] = aten::split(%x.1, %7, %8)
            return ()
        """
        graph = parse_ir(graph_str)
        alias_db = graph.alias_db()
        split_node = graph.findNode("aten::split")
        # split input enters wildcard set, list initalized as containing wildcard set
        self.assertTrue(
            alias_db.may_contain_alias(next(split_node.inputs()), split_node.output())
        )
        # because %x.1 enters wildcard set, it now aliases other members of wildcard set (graph inputs)
        self.assertTrue(
            alias_db.may_contain_alias(next(split_node.inputs()), next(graph.inputs()))
        )

    def test_nested_list_construct_not_wildcard(self):
        @torch.jit.script
        def foo(x):
            y = torch.rand([2, 2])
            return [y]

        graph = foo.graph
        graph.alias_db()
        alias_db = graph.alias_db()
        ten_construct = graph.findNode("aten::rand").output()
        output = next(graph.outputs())
        self.assertTrue(alias_db.may_contain_alias(ten_construct, output))
        self.assertFalse(
            alias_db.may_contain_alias(next(graph.inputs()), ten_construct)
        )

    def test_recursive_calls(self):
        @torch.jit.script
        def foo(x, y):
            x.add_(1)
            return x + y

        @torch.jit.script
        def caller():
            a = torch.rand([2, 2])
            b = torch.ones([2, 2])
            out1 = foo(a, b)
            c = torch.rand([1])
            d = torch.ones([2])
            out2 = foo(d, c)
            return out1, out2

        isFrozen = False
        descend_function_calls = True
        alias_db = caller.graph.alias_db(isFrozen, descend_function_calls)
        func_calls = caller.graph.findAllNodes("prim::CallFunction")
        self.assertEqual(len(func_calls), 2)
        for node in func_calls:
            inps = list(node.inputs())
            self.assertTrue(alias_db.has_writers(inps[1]))
            self.assertFalse(alias_db.has_writers(inps[2]))

        class Mod(torch.nn.Module):
            def forward(self):
                a = torch.rand([2, 2])
                b = torch.ones([2, 2])
                out1 = self.foo2(a, b)
                c = torch.rand([1])
                d = torch.ones([2])
                out2 = self.foo2(d, c)
                return out1, out2

            def foo2(self, x, y):
                x.add_(1)
                return x + y

        mod = torch.jit.script(Mod())
        alias_db = mod.graph.alias_db(isFrozen, descend_function_calls)
        func_calls = mod.graph.findAllNodes("prim::CallMethod")
        self.assertEqual(len(func_calls), 2)
        for node in func_calls:
            inps = list(node.inputs())
            self.assertTrue(alias_db.has_writers(inps[1]))
            self.assertFalse(alias_db.has_writers(inps[2]))

    def test_multiple_compilation_units(self):
        # This is a repro of an internal issue we saw.
        # Here, we have a large number (40) of modules each with the same name (MyModuleCUTest).
        # AliasDB uses some hash tables that hash on types; each of these 40 modules are not
        # identical because they have different compilation units, but they have the same name.
        # Therefore, if we hash only on the module name (which we previously did), we will have
        # hash collisions for all of these module types.
        #
        # flat_hash_map has very bad performance (exponential) for this hash collision behavior.
        # This OOMs prior to the fix.
        N = 40

        class MultiTmpFile:
            def __init__(self, N):
                self.N = N
                self.ctxs = [
                    TemporaryFileName(mode="w", suffix=".py") for _ in range(N)
                ]

            def __enter__(self):
                return [x.__enter__() for x in self.ctxs]

            def __exit__(self, exc_type, exc_value, traceback):
                return [x.__exit__(exc_type, exc_value, traceback) for x in self.ctxs]

        class ModuleWrapper(torch.nn.Module):
            def __init__(self, module_list):
                super().__init__()
                self.module_list = module_list

            def forward(self, x):
                for mod in self.module_list:
                    x = mod(x)
                return x

        with MultiTmpFile(N) as fnames:
            module_list = torch.nn.ModuleList()
            global MyModuleCUTest

            class MyModuleCUTest(torch.nn.Module):
                def forward(self, x):
                    return x + 2

            for _, fname in enumerate(fnames):
                mod = torch.jit.script(MyModuleCUTest())
                torch.jit.save(mod, fname)
                loaded_mod = torch.jit.load(fname)
                module_list.append(loaded_mod)

            mod = ModuleWrapper(module_list)
            mod = torch.jit.script(mod)
            mod(torch.zeros((2, 2)))