test_peephole.py

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

import unittest
from typing import Callable, List

import torch
from torch import nn
from torch.testing import FileCheck
from torch.testing._internal.jit_utils import _inline_everything, JitTestCase, RUN_CUDA

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 TestPeephole(JitTestCase):
    def test_peephole_with_writes(self):
        def test_write(x):
            s = 0
            s += x
            s += x
            return s

        self.checkScript(test_write, (torch.ones(4, 4),))

    def test_peephole_with_non_output_writes(self):
        @torch.jit.ignore
        def nomnom(x):
            pass

        def test_write(x):
            t = torch.ones_like(x)
            z = x.clone()
            y = z + 0
            z.add_(t)
            # this makes sure z isn't blasted out of existence
            # because it isn't returned or used in a side-effectful
            # way
            nomnom(z)
            return y + y

        a = torch.ones(4, 4)
        j = self.checkScript(test_write, (a,))

    def test_peephole_no_output_aliasing(self):
        def test_peephole(x):
            y = x + 0
            return x, y

        a = torch.ones(4, 4)
        j = self.checkScript(test_peephole, (a,))
        r1, r2 = j(a)
        self.assertNotEqual(r1.data_ptr(), r2.data_ptr())

    def test_peephole(self):
        a = torch.tensor([0.4])
        b = torch.tensor([0.7])
        c = torch.tensor([0], dtype=torch.int32)

        def f(x, y):
            return x.type_as(y)

        tf = torch.jit.trace(f, (a, b))
        FileCheck().check("type_as").run(str(tf.graph))
        self.run_pass("peephole", tf.graph)
        FileCheck().check_not("type_as").run(str(tf.graph))
        tf2 = torch.jit.trace(f, (a, c))
        s = str(tf2.graph)
        self.run_pass("peephole", tf2.graph)
        self.assertEqual(s, str(s))

    def test_peephole_dynamic(self):
        def f(x, y):
            return x.type_as(y)

        fn = torch.jit.script(f)
        s = str(fn.graph)
        torch._C._jit_pass_peephole(fn.graph)
        self.assertEqual(s, str(fn.graph))

    def test_peephole_list_ops(self):
        @torch.jit.script
        def foo(x, y, z):
            return len([x, y, z])

        self.run_pass("peephole", foo.graph)
        FileCheck().check("value=3").check_next("return").run(foo.graph)

        @torch.jit.script
        def foo(x, y, z):
            li = [x, y, z]
            for i in range(len(x)):
                li.append(x)
            return len([x, y, z])

        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("aten::len").run(foo.graph)

        @torch.jit.script
        def foo(x, y, z):
            li = [x, y, z]
            return li[1], li[-2]

        FileCheck().check("aten::__getitem__").run(foo.graph)
        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("aten::__getitem__").run(foo.graph)

        @torch.jit.script
        def foo(x, y, z):
            li = [x, y, z]
            return li[-7]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("aten::__getitem__").run(foo.graph)

        @torch.jit.script
        def foo(x, y, z):
            li = [x, y, z]
            for i in range(len(x)):
                li.append(x)
            return li[-2]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("aten::__getitem__").run(foo.graph)

    @unittest.skipIf(not RUN_CUDA, "cpp tests require CUDA")
    def test_peephole_cuda(self):
        a = torch.tensor([0.4], device="cpu")
        b = torch.tensor([0.7], device="cuda")
        c = torch.tensor([0.7], device="cuda")

        def f(x, y):
            return x.type_as(y)

        trace = torch.jit.trace(f, (a, c))
        s = str(trace.graph)
        self.run_pass("peephole", trace.graph)
        self.assertEqual(s, str(trace.graph))
        trace = torch.jit.trace(f, (b, c))
        self.run_pass("peephole", trace.graph)
        self.run_pass("dce", trace.graph)
        FileCheck().check_not("type_as").run(str(trace.graph))

    @_inline_everything
    def test_peephole_type_refinements(self):
        def refine(x):
            # type: (Optional[Tensor]) -> Tensor
            return x if x is not None else torch.tensor(3)

        @torch.jit.script
        def test():
            return refine(torch.tensor(4))

        FileCheck().check("prim::unchecked_cast").run(test.graph)
        self.run_pass("peephole", test.graph)
        FileCheck().check_not("prim::unchecked_cast").run(test.graph)

        # refinement not optimzied out
        def is_int_tensor(x):
            scalar = x.item()
            if isinstance(scalar, int):
                return scalar + 3
            else:
                return 8

        self.checkScript(is_int_tensor, (torch.tensor(2),))
        self.checkScript(is_int_tensor, (torch.tensor(2.5),))
        graph = torch.jit.script(is_int_tensor).graph
        self.run_pass("peephole", graph)
        FileCheck().check("prim::unchecked_cast").run(graph)

    def test_short_circuit_optimization(self):
        @torch.jit.script
        def const_expressions(x):
            # type: (int) -> Tuple[bool, bool]
            return x == 1 and False, x == 1 or True

        self.run_pass("constant_propagation", const_expressions.graph)
        FileCheck().check_not("prim::If").check_not("aten::eq").run(
            const_expressions.graph
        )
        self.assertEqual(const_expressions(1), (False, True))

        @torch.jit.script
        def redundant_expressions(x):
            # type: (int) -> Tuple[bool, bool]
            return x == 1 and True, x == 1 or False

        self.run_pass("peephole", redundant_expressions.graph)
        self.assertEqual(redundant_expressions(1), (True, True))
        self.assertEqual(redundant_expressions(0), (False, False))
        # and True / or False are removed from graph
        FileCheck().check("aten::eq").check_not("prim::If").run(
            redundant_expressions.graph
        )

    def test_conv_dim_folding(self):
        modules = [nn.Conv1d, nn.Conv2d, nn.Conv3d]
        for mod in modules:

            class ConvDim(torch.nn.Module):
                def __init__(self):
                    super().__init__()
                    self.conv = mod(3, 32, kernel_size=3, stride=2, bias=False)

                def forward(self, x):
                    x = self.conv(x)
                    return x.dim()

            conv_dim = torch.jit.script(ConvDim())
            self.run_pass("inline", conv_dim.graph)
            self.run_pass("peephole", conv_dim.graph)
            FileCheck().check_not("conv").check_not("dim").run(conv_dim.graph)

            class ConvDimMutate(torch.nn.Module):
                def __init__(self):
                    super().__init__()
                    self.conv = mod(3, 32, kernel_size=3, stride=2, bias=False)

                def forward(self, x):
                    x = self.conv(x)
                    x.resize_([4, 4])
                    return x.dim()

            conv_dim = torch.jit.script(ConvDimMutate())
            self.run_pass("inline", conv_dim.graph)
            self.run_pass("peephole", conv_dim.graph)
            FileCheck().check("conv").check("dim").run(conv_dim.graph)

    def test_normalized_rsub(self):
        a = torch.tensor([1, 2, 3])
        b = torch.tensor([4, 5, 6])

        def convertible_rsub(x, y):
            return (x - y), torch.rsub(y, x)

        self.checkScript(convertible_rsub, (a, b))
        op_graph = torch.jit.script(convertible_rsub).graph
        FileCheck().check_count("aten::sub", 2, exactly=True).run(op_graph)
        FileCheck().check_count("aten::rsub", 0, exactly=True).run(op_graph)

    def test_normalized_is_op(self):
        def convertible_is_op(x: bool, y: bool):
            return x is True, False is x, x is y

        self.checkScript(convertible_is_op, (True, False))

        op_graph = torch.jit.script(convertible_is_op).graph
        FileCheck().check_count("aten::eq", 3, exactly=True).run(op_graph)
        FileCheck().check_count("aten::__is__", 0, exactly=True).run(op_graph)

    def test_normalized_isnot_op(self):
        def convertible_isnot_op(x: bool, y: bool):
            return x is not True, False is not x, x is not y

        self.checkScript(convertible_isnot_op, (True, False))

        op_graph = torch.jit.script(convertible_isnot_op).graph
        FileCheck().check_count("aten::ne", 3, exactly=True).run(op_graph)
        FileCheck().check_count("aten::__isnot__", 0, exactly=True).run(op_graph)

    def test_peephole_list_len(self):
        def run_peephole_and_check_const_value(graph, const_string):
            torch._C._jit_pass_peephole_list_idioms(graph, refine_list_len=True)
            self.run_pass("constant_propagation", graph)
            FileCheck().check(const_string).check_next("return").run(graph)

        def gen_li(inp_len: int):
            return [0 for i in range(inp_len)]

        @torch.jit.script
        def foo(x: List[int], y: List[int]):
            if len(x) != 4 or len(y) != 5:
                raise Exception("")  # noqa: TRY002

            return len(x) + len(y)

        run_peephole_and_check_const_value(foo.graph, "value=9")
        self.assertEqual(foo(gen_li(4), gen_li(5)), 9)
        with self.assertRaises(Exception):
            foo(2, 4)

        @torch.jit.script
        def foo(x: List[int], y: List[int]):
            if len(x) == 4 and len(y) == 5:
                pass
            else:
                raise Exception("hi")  # noqa: TRY002

            return len(x) + len(y)

        run_peephole_and_check_const_value(foo.graph, "value=9")
        self.assertEqual(foo(gen_li(4), gen_li(5)), 9)
        with self.assertRaises(Exception):
            foo(2, 4)

        @torch.jit.script
        def foo(x: List[int], y: List[int], z: List[int]):
            if len(x) != 4:
                raise Exception("..")  # noqa: TRY002
            else:
                if len(y) != 8:
                    raise Exception("...")  # noqa: TRY002
                else:
                    if len(z) == 3:
                        pass
                    else:
                        raise Exception("...")  # noqa: TRY002

            return len(x) + len(y) * len(z)

        run_peephole_and_check_const_value(foo.graph, "value=28")
        self.assertEqual(foo(gen_li(4), gen_li(8), gen_li(3)), 28)
        with self.assertRaises(Exception):
            foo(1, 2, 3)

        # refinement should persist in second len(x) call

        @torch.jit.script
        def foo(x: List[int], cond: bool):
            if len(x) == 4:
                if cond:
                    return len(x)
                return 4

            return 4

        run_peephole_and_check_const_value(foo.graph, "value=4")

        def test_const_tuple_output(graph, const_inputs):
            tup = graph.findNode("prim::TupleConstruct")
            for i, elem in enumerate(tup.inputs()):
                if i in const_inputs:
                    self.assertIsNotNone(elem.toIValue())
                else:
                    self.assertIsNone(elem.toIValue())

        # testing combinations of x1 : {True, False} x
        # {then/else branch} x assert {True/False}

        @torch.jit.script
        def foo(x: List[int], b: List[int]):
            if len(x) == 5:
                x1 = True
            else:
                x1 = len(b) != 4
            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
            return len(x), len(b)

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        torch._C._jit_pass_constant_propagation(foo.graph)
        # we can only infer len(b) == 4 here
        test_const_tuple_output(foo.graph, [1])

        @torch.jit.script
        def foo(x: List[int], b: List[int]):
            if len(x) == 5:
                x1 = False
            else:
                x1 = len(b) != 4
            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
            return len(x), len(b)

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        torch._C._jit_pass_constant_propagation(foo.graph)
        # cant infer anything
        test_const_tuple_output(foo.graph, [])

        @torch.jit.script
        def foo(x: List[int], b: List[int]):
            if len(x) == 5:
                x1 = True
            else:
                x1 = len(b) == 4
            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
            return len(x), len(b)

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        torch._C._jit_pass_constant_propagation(foo.graph)
        # we cant infer anything, only len(b) != 4
        test_const_tuple_output(foo.graph, [])

        @torch.jit.script
        def foo(x: List[int], b: List[int]):
            if len(x) == 5:
                x1 = True
            else:
                x1 = len(b) != 4
            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
            return len(x), len(b)

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        torch._C._jit_pass_constant_propagation(foo.graph)
        # can infer len(b) == 4
        test_const_tuple_output(foo.graph, [1])

        # swap branches
        @torch.jit.script
        def foo(x: List[int], b: List[int]):
            if len(x) != 5:
                x1 = len(b) != 4
            else:
                x1 = True
            assert x1 == False  # noqa: E712 TODO: canonicalize x is False to aten::eq
            return len(x), len(b)

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        torch._C._jit_pass_constant_propagation(foo.graph)
        # can infer len(b) == 4
        test_const_tuple_output(foo.graph, [1])

        # use __not__
        @torch.jit.script
        def foo(x: List[int], b: List[int]):
            if len(x) != 5:
                x1 = len(b) != 4
            else:
                x1 = True
            assert not x1
            return len(x), len(b)

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        torch._C._jit_pass_constant_propagation(foo.graph)
        # can infer len(b) == 4
        test_const_tuple_output(foo.graph, [1])

        # Test unsuccessful optimizations

        @torch.jit.script
        def foo(x: List[int]):
            assert len(x) == 4
            x.append(3)
            return len(x)

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        self.run_pass("constant_propagation", foo.graph)
        FileCheck().check_count("aten::len", 2).run(foo.graph)

        @torch.jit.script
        def foo(x: List[int], y: List[int]):
            assert len(x) == 4 or len(y) == 5
            return len(x) + len(y)

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        self.run_pass("constant_propagation", foo.graph)
        FileCheck().check_count("aten::len", 4).run(foo.graph)

    def test_integer_refinement(self):
        def run_peephole_and_check_const_value(graph, const_string):
            self.run_pass("refine_integer_values", graph)
            self.run_pass("constant_propagation", graph)
            self.run_pass("dce", graph)
            FileCheck().check(const_string).check_next("return").run(graph)

        @torch.jit.script
        def foo(x: int, y: int):
            if x != 4 or y != 5:
                raise Exception("")  # noqa: TRY002

            return x + y

        graph = foo.graph
        self.run_pass("refine_integer_values", graph)
        self.run_pass("constant_propagation", graph)
        self.run_pass("dce", graph)

        run_peephole_and_check_const_value(foo.graph, "value=9")
        self.assertEqual(foo(4, 5), 9)
        with self.assertRaises(Exception):
            foo(2, 4)

        @torch.jit.script
        def foo(x: int, y: int):
            if x == 4 and y == 5:
                pass
            else:
                raise Exception("hi")  # noqa: TRY002

            return x + y

        run_peephole_and_check_const_value(foo.graph, "value=9")
        self.assertEqual(foo(4, 5), 9)
        with self.assertRaises(Exception):
            foo(2, 4)

        @torch.jit.script
        def foo(x: int, y: int, z: int):
            if x != 4:
                raise Exception("..")  # noqa: TRY002
            else:
                if y != 8:
                    raise Exception("...")  # noqa: TRY002
                else:
                    if z == 3:
                        pass
                    else:
                        raise Exception("...")  # noqa: TRY002

            return x + y * z

        run_peephole_and_check_const_value(foo.graph, "value=28")
        self.assertEqual(foo(4, 8, 3), 28)
        with self.assertRaises(Exception):
            foo(1, 2, 3)

        # refinement should persist in second len(x) call

        @torch.jit.script
        def foo(x: int, cond: bool):
            if x == 4:
                if cond:
                    return x
                return 4

            return 4

        run_peephole_and_check_const_value(foo.graph, "value=4")

        @torch.jit.script
        def foo(x: int, y: int):
            assert x == 4 or y == 5
            return x + y

        torch._C._jit_pass_peephole_list_idioms(foo.graph, refine_list_len=True)
        self.run_pass("constant_propagation", foo.graph)
        FileCheck().check("aten::add").run(foo.graph)

    def test_optimize_out_comparison_same_value(self):
        def foo(x: int):
            return x == x, x != x

        def foo2(x: List[int]):
            return x == x, x != x

        for func, inp in zip([foo, foo2], [1, [2, 3]]):
            func_s = torch.jit.script(func)
            self.run_pass("peephole", func_s.graph)
            FileCheck().check_not("aten::eq").check_not("aten::neq").run(func_s.graph)
            self.assertEqual(func(inp), func_s(inp))

    def test_peephole_add_zero(self):
        @torch.jit.script
        def foo(x: int):
            return x + 0, 0 + x

        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("aten::add")
        self.assertEqual(foo(3), (3, 3))

    def test_noop_peephole(self):
        # test unsuccessful
        def foo1(x):
            return x + 0

        def foo2():
            x = torch.zeros([2, 2])
            x.sub_(3)
            return x + 0

        def foo3():
            x = torch.zeros([2, 2])
            return x, x + 0

        def foo4():
            x = torch.zeros([2, 2])
            return x + 0.0

        funcs = foo1, foo2, foo3, foo4
        inps = (torch.ones([2]),), (), (), ()
        for func, inp in zip(funcs, inps):
            foo_s = torch.jit.script(func)
            self.run_pass("peephole", foo_s.graph)
            FileCheck().check_count("aten::add", 1, exactly=True).run(foo_s.graph)
            self.assertEqual(func(*inp), foo_s(*inp))

        # successful
        def func(x):
            return (x + 0) * 1 - 5

        func_s = torch.jit.script(func)
        self.run_pass("peephole", func_s.graph)
        # bail on modified value first
        FileCheck().check_not("aten::add").check("aten::mul").run(func_s.graph)
        # second run it should succeed
        self.run_pass("peephole", func_s.graph)
        FileCheck().check_not("aten::add").check_not("aten::mul").run(func_s.graph)
        self.assertEqual(func(torch.ones([2, 2])), func_s(torch.ones([2, 2])))

        def func(x):
            return (x + 0.0) - 5

        func_s = torch.jit.script(func)
        inp = next(func_s.graph.inputs())
        inp.setType(torch._C.TensorType.create_from_tensor(torch.rand([2, 2])))
        torch._C._jit_pass_peephole(func_s.graph, disable_shape_peepholes=True)
        FileCheck().check("aten::add").run(func_s.graph)
        torch._C._jit_pass_peephole(func_s.graph, disable_shape_peepholes=False)
        FileCheck().check_not("aten::add").run(func_s.graph)

    def test_refine_integer_values(self):
        @torch.jit.script
        def foo(x: int):
            y = 1
            if x == 1:
                return y
            else:
                return x

        self.run_pass("refine_integer_values", foo.graph)
        self.run_pass("constant_propagation", foo.graph)
        self.run_pass("dce", foo.graph)
        FileCheck().check("graph").check_next("return").run(foo.graph)
        self.assertEqual(foo(2), 2)
        self.assertEqual(foo(1), 1)

    def test_peephole_len_list(self):
        @torch.jit.script
        def foo(x):
            return len(x.size())

        self.run_pass("peephole", foo.graph)
        FileCheck().check("aten::len").run(foo.graph)
        inputs = list(foo.graph.inputs())
        inputs[0].setType(inputs[0].type().with_sizes([None, None]))
        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("aten::len").run(foo.graph)
        self.assertEqual(2, foo(torch.rand([3, 1])))

        @torch.jit.script
        def foo(x):
            li = x.size()
            li.append(4)
            return len(li)

        inputs = list(foo.graph.inputs())
        inputs[0].setType(inputs[0].type().with_sizes([None, None]))
        self.run_pass("peephole", foo.graph)
        FileCheck().check("aten::len").run(foo.graph)
        self.assertEqual(3, foo(torch.rand([3, 1])))

    def test_peephole_optional_refine(self):
        @torch.jit.script
        def foo(z: int, z2: int, cond: bool):
            if cond:
                return z
            else:
                return z2

        out = next(foo.graph.findNode("prim::If").outputs())
        out.setType(torch._C.OptionalType(torch._C.IntType.get()))
        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("int?").run(foo.graph)

    def test_peephole_int(self):
        @torch.jit.script
        def foo(x):
            # type: (number)
            return int(x)

        FileCheck().check("aten::Int").run(foo.graph)
        next(foo.graph.inputs()).setType(torch._C.IntType.get())
        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("aten::Int").run(foo.graph)

    def test_peephole_arith(self):
        @torch.jit.script
        def foo(input0: int, input1: int, input2: int, input3: int):
            _1 = torch.add(input1, 2)
            _3 = torch.add(input3, 2)
            _5 = torch.add(1, torch.sub(_1, 3) // 1)
            _6 = torch.add(1 * torch.sub(_3, 3) // 1, 1) / 1
            return [_5, int(_6)]

        FileCheck().check("aten::add").check("aten::sub").check("aten::mul").check(
            "aten::floordiv"
        ).check("aten::div").run(foo.graph)
        self.run_pass("peephole", foo.graph)
        FileCheck().check("graph").check("):").check_next("ListConstruct").check_next(
            "return"
        ).run(foo.graph)
        self.assertEqual(foo(0, 1, 2, 3), [1, 3])

    def test_peephole_dict_getitem_simple(self):
        @torch.jit.script
        def foo(a: int, b: int):
            d = {0: a, 1: b}
            x = d[1]
            y = d[0]
            return x, y

        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("DictConstruct").check_not("__getitem__").run(foo.graph)
        self.assertEqual(foo(0, 1), (1, 0))

        @torch.jit.script
        def foo(a: int, b: int):
            d = {"0": a, "1": b}
            x = d["1"]
            y = d["0"]
            return x, y

        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("DictConstruct").check_not("__getitem__").run(foo.graph)
        self.assertEqual(foo(0, 1), (1, 0))

        @torch.jit.script
        def foo(a: int, b: int):
            d = {0.0: a, 1.0: b}
            x = d[1.0]
            y = d[0.0]
            return x, y

        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("DictConstruct").check_not("__getitem__").run(foo.graph)
        self.assertEqual(foo(0, 1), (1, 0))

    def test_peephole_dict_getitem_no_optimization_missing_key(self):
        @torch.jit.script
        def foo():
            d = {0: 1}
            return d[2]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("__getitem__").run(foo.graph)

    def test_peephole_dict_getitem_no_optimization_get_input_arg(self):
        # Here we don't know if the input arg is in the dict, so we can't
        # make the optimization.
        @torch.jit.script
        def foo(a: int):
            d = {0: 1}
            return d[a]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("__getitem__").run(foo.graph)
        self.assertEqual(foo(0), 1)

    def test_peephole_dict_getitem_no_optimization_dict_modified(self):
        @torch.jit.script
        def foo():
            d = {0: 1}
            d[0] = 2
            return d[0]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("__getitem__").run(foo.graph)
        self.assertEqual(foo(), 2)

    def test_peephole_dict_getitem_no_optimization_overlapping_keys(self):
        @torch.jit.script
        def foo():
            d = {0: 1, 0: 2}  # noqa: F601
            return d[0]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("__getitem__").run(foo.graph)

    def test_peephole_dict_getitem_no_optimization_keys_might_overlap(self):
        @torch.jit.script
        def foo(x: int):
            d = {0: 1, x: 2}
            return d[x]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("__getitem__").run(foo.graph)

    def test_peephole_dict_getitem_no_optimization_unsupported_type(self):
        @torch.jit.script
        def foo():
            a = torch.rand((2, 2))
            d = {a: 1}
            return d[a]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("__getitem__").run(foo.graph)
        self.assertEqual(foo(), 1)

    def test_peephole_dict_len(self):
        @torch.jit.script
        def foo():
            d = {0: 1, 1: 2}
            return len(d)

        self.run_pass("peephole", foo.graph)
        FileCheck().check_not("DictConstruct").check_not("len").run(foo.graph)
        self.assertEqual(foo(), 2)

    def test_peephole_dict_len_no_optimization_overlapping_keys(self):
        @torch.jit.script
        def foo():
            d = {0: 1, 0: 2}  # noqa: F601
            return len(d)

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("len").run(foo.graph)
        self.assertEqual(foo(), 1)

    def test_peephole_dict_len_no_optimization_keys_might_overlap(self):
        @torch.jit.script
        def foo(x: int):
            d = {0: 1, x: 2}
            return len(d)

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("len").run(foo.graph)

    def test_peephole_dict_len_no_optimization_unsupported_type(self):
        @torch.jit.script
        def foo():
            a = torch.rand((2, 2))
            d = {a: 1}
            return len(d)

        self.run_pass("peephole", foo.graph)
        FileCheck().check("DictConstruct").check("len").run(foo.graph)
        self.assertEqual(foo(), 1)

    def test_peephole_slice_all_three_args(self):
        def foo(x: int):
            return [1, 2, x, 4, 5, 6, 7][-5:6:2]

        graph = torch.jit.script(foo).graph
        self.run_pass("peephole", graph)
        FileCheck().check_not("aten::slice").run(graph)
        self.checkScript(foo, (3,))

    def test_peephole_slice_one_empty_arg(self):
        def check_helper(fn: Callable[[int], None]) -> None:
            graph = torch.jit.script(fn).graph
            self.run_pass("peephole", graph)
            FileCheck().check_not("aten::slice").run(graph)
            self.checkScript(fn, (3,))

        def foo(x: int):
            return [1, 2, x, 4, 5, 6, 7][1::2]

        check_helper(foo)

        def foo(x: int):
            return [1, 2, x, 4, 5, 6, 7][:5:3]

        check_helper(foo)

        def foo(x: int):
            return [1, 2, x, 4, 5, 6, 7][0:4]

        check_helper(foo)

    def test_peephole_slice_two_empty_args(self):
        def check_helper(fn: Callable[[int], None]) -> None:
            graph = torch.jit.script(fn).graph
            self.run_pass("peephole", graph)
            FileCheck().check_not("aten::slice").run(graph)
            self.checkScript(fn, (3,))

        def foo(x: int):
            return [1, 2, x, 4, 5, 6, 7][::2]

        check_helper(foo)

        def foo(x: int):
            return [1, 2, x, 4, 5, 6, 7][:5]

        check_helper(foo)

        def foo(x: int):
            return [1, 2, x, 4, 5, 6, 7][1:]

        check_helper(foo)

    def test_peephole_slice_optimization_not_applied_list_modified(self):
        @torch.jit.script
        def foo():
            li = [1, 2, 3, 4, 5, 6, 7]
            li[0] = 0
            return li[2:5]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("aten::slice").run(foo.graph)

    def test_peephole_slice_optimization_not_applied_non_const_args(self):
        @torch.jit.script
        def foo(x: int, y: int):
            li = [1, 2, 3, 4, 5, 6, 7]
            return li[x:y]

        self.run_pass("peephole", foo.graph)
        FileCheck().check("aten::slice").run(foo.graph)