Set pyre-strict for passes unit tests.

Differential Revision: D76767018

Author: eigen-k

backends/cadence/aot/tests/test_fusion_ops_passes.py

@@ -4,11 +4,11 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# pyre-unsafe
+# pyre-strict
 
 import math
 import unittest
-from typing import cast, Optional
+from typing import cast, List, Optional
 
 import executorch.backends.cadence.aot.ops_registrations  # noqa
 import torch
@@ -224,11 +224,11 @@ def verify_nop_memory_alloc(self, graph_module: torch.fx.GraphModule) -> None:
     # GenerateSliceAndSelectNopConstraints, and GenerateCatNopConstraints passes.
     def run_memory_planning(
         self,
-        original,
-        opt_level=2,
-        mem_algo=1,  # greedy_by_size_for_offset_calculation_with_hierarchy
-        alloc_graph_input=True,
-        alloc_graph_output=True,
+        original: GraphModule,
+        opt_level: int = 2,
+        mem_algo: int = 1,  # greedy_by_size_for_offset_calculation_with_hierarchy
+        alloc_graph_input: bool = True,
+        alloc_graph_output: bool = True,
         memory_config: Optional[MemoryConfig] = None,
     ) -> GraphModule:
         if memory_config is None:
@@ -242,6 +242,7 @@ def run_memory_planning(
             alloc_graph_output=alloc_graph_output,
         )(graph_module).graph_module
 
+    # pyre-ignore[56]
     @parameterized.expand(
         [
             [
@@ -259,7 +260,11 @@ def run_memory_planning(
         ]
     )
     def test_optimize_cat_on_placeholders(
-        self, x_shape, y_shape, concat_dim, alloc_graph_input
+        self,
+        x_shape: List[int],
+        y_shape: List[int],
+        concat_dim: int,
+        alloc_graph_input: bool,
     ) -> None:
         concat_shape = [x_shape[concat_dim] + y_shape[concat_dim], x_shape[1]]
         builder = GraphBuilder()
@@ -294,7 +299,12 @@ def test_optimize_cat_on_placeholders(
     # "add_add_cat_model" : cat(x + 123, y + 456)
     # "add_add_cat_add_model": cat(x + 123, y + 456) + 789
     def get_graph_module(
-        self, model_name, x_shape, y_shape, concated_shape, concat_dim
+        self,
+        model_name: str,
+        x_shape: List[int],
+        y_shape: List[int],
+        concated_shape: List[int],
+        concat_dim: int,
     ) -> GraphModule:
         builder = GraphBuilder()
         x = builder.placeholder("x", torch.ones(*x_shape, dtype=torch.float32))
@@ -346,6 +356,7 @@ def get_graph_module(
 
         raise ValueError(f"Unknown model name {model_name}")
 
+    # pyre-ignore[56]
     @parameterized.expand(
         [
             (
@@ -366,7 +377,12 @@ def get_graph_module(
         name_func=lambda f, _, param: f"{f.__name__}_{param.args[0]}",
     )
     def test_cat_optimized(
-        self, _, x_shape, y_shape, concated_shape, concat_dim
+        self,
+        _,
+        x_shape: List[int],
+        y_shape: List[int],
+        concated_shape: List[int],
+        concat_dim: int,
     ) -> None:
         original = self.get_graph_module(
             "add_add_cat_model", x_shape, y_shape, concated_shape, concat_dim
@@ -379,6 +395,7 @@ def test_cat_optimized(
         self.assertEqual(count_node(graph_module, torch.ops.aten._cat_nop.out), 1)
         self.verify_nop_memory_alloc(graph_module)
 
+    # pyre-ignore[56]
     @parameterized.expand(
         [
             (
@@ -392,7 +409,12 @@ def test_cat_optimized(
         name_func=lambda f, _, param: f"{f.__name__}_{param.args[0]}",
     )
     def test_cat_not_optimized(
-        self, _, x_shape, y_shape, concated_shape, concat_dim
+        self,
+        _,
+        x_shape: List[int],
+        y_shape: List[int],
+        concated_shape: List[int],
+        concat_dim: int,
     ) -> None:
         original = self.get_graph_module(
             "add_add_cat_model", x_shape, y_shape, concated_shape, concat_dim
@@ -404,6 +426,7 @@ def test_cat_not_optimized(
         self.assertEqual(count_node(graph_module, torch.ops.aten.cat.out), 1)
         self.verify_nop_memory_alloc(graph_module)
 
+    # pyre-ignore[56]
     @parameterized.expand(
         [
             (
@@ -426,7 +449,13 @@ def test_cat_not_optimized(
         name_func=lambda f, _, param: f"{f.__name__}_{param.args[0]}",
     )
     def test_cat_not_graph_output(
-        self, _, x_shape, y_shape, concated_shape, concat_dim, expected_cat_nodes
+        self,
+        _,
+        x_shape: List[int],
+        y_shape: List[int],
+        concated_shape: List[int],
+        concat_dim: int,
+        expected_cat_nodes: int,
     ) -> None:
         original = self.get_graph_module(
             "add_add_cat_add_model", x_shape, y_shape, concated_shape, concat_dim
@@ -493,13 +522,14 @@ def test_optimize_cat_with_slice(self) -> None:
         self.assertEqual(count_node(graph_module, exir_ops.edge.aten.slice.Tensor), 1)
         self.verify_nop_memory_alloc(graph_module)
 
+    # pyre-ignore[56]
     @parameterized.expand(
         [
             (True,),  # alloc_graph_input
             (False,),  # alloc_graph_input
         ],
     )
-    def test_optimize_cat_with_slice_infeasible(self, alloc_graph_input) -> None:
+    def test_optimize_cat_with_slice_infeasible(self, alloc_graph_input: bool) -> None:
         x_shape = [5, 6]
         y_shape = [3, 6]
         concated_shape = [8, 6]

backends/cadence/aot/tests/test_memory_passes.py

@@ -4,13 +4,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# pyre-unsafe
+# pyre-strict
 
 
 import unittest
-
 from copy import deepcopy
 
+from typing import Callable, Dict
+
 from executorch.backends.cadence.aot import pass_utils
 from executorch.backends.cadence.aot.pass_utils import (
     ALL_CADENCE_PASSES,
@@ -23,24 +24,26 @@
 
 
 class TestBase(unittest.TestCase):
-    def setUp(self):
+    def setUp(self) -> None:
         # Before running each test, create a copy of _all_passes to later restore it after test.
         # This avoids messing up the original _all_passes when running tests.
         self._all_passes_original = deepcopy(ALL_CADENCE_PASSES)
         # Clear _all_passes to do a clean test. It'll be restored after each test in tearDown().
         pass_utils.ALL_CADENCE_PASSES.clear()
 
-    def tearDown(self):
+    def tearDown(self) -> None:
         # Restore _all_passes to original state before test.
         pass_utils.ALL_CADENCE_PASSES = self._all_passes_original
 
-    def get_filtered_passes(self, filter_):
+    def get_filtered_passes(
+        self, filter_: Callable[[ExportPass], bool]
+    ) -> Dict[ExportPass, CadencePassAttribute]:
         return {cls: attr for cls, attr in ALL_CADENCE_PASSES.items() if filter_(cls)}
 
 
 # Test pass registration
 class TestPassRegistration(TestBase):
-    def test_register_cadence_pass(self):
+    def test_register_cadence_pass(self) -> None:
         pass_attr_O0 = CadencePassAttribute(opt_level=0)
         pass_attr_debug = CadencePassAttribute(opt_level=None, debug_pass=True)
         pass_attr_O1_all_backends = CadencePassAttribute(
@@ -73,7 +76,7 @@ class DummyPass_Debug(ExportPass):
 
 # Test pass filtering
 class TestPassFiltering(TestBase):
-    def test_filter_none(self):
+    def test_filter_none(self) -> None:
         pass_attr_O0 = CadencePassAttribute(opt_level=0)
         pass_attr_O1_debug = CadencePassAttribute(opt_level=1, debug_pass=True)
         pass_attr_O1_all_backends = CadencePassAttribute(
@@ -103,7 +106,7 @@ class DummyPass_O1_All_Backends(ExportPass):
         }
         self.assertEqual(O1_filter_passes, expected_passes)
 
-    def test_filter_debug(self):
+    def test_filter_debug(self) -> None:
         pass_attr_O1_debug = CadencePassAttribute(opt_level=1, debug_pass=True)
         pass_attr_O2 = CadencePassAttribute(opt_level=2)
 
@@ -122,7 +125,7 @@ class DummyPass_O2(ExportPass):
         # chooses debug=False.
         self.assertEqual(debug_filter_passes, {DummyPass_O2: pass_attr_O2})
 
-    def test_filter_all(self):
+    def test_filter_all(self) -> None:
         @register_cadence_pass(CadencePassAttribute(opt_level=1))
         class DummyPass_O1(ExportPass):
             pass
@@ -138,7 +141,7 @@ class DummyPass_O2(ExportPass):
         # passes with opt_level <= 0
         self.assertEqual(debug_filter_passes, {})
 
-    def test_filter_opt_level_None(self):
+    def test_filter_opt_level_None(self) -> None:
         pass_attr_O1 = CadencePassAttribute(opt_level=1)
         pass_attr_O2_debug = CadencePassAttribute(opt_level=2, debug_pass=True)