Remove pass-local ExportPass replays (#22156)

backends/cadence/aot/fuse_ops.py

@@ -712,6 +712,16 @@ def get_fused_node(
                 out_dtype=out_dtype,
                 graph=graph_module.graph,
             )
+            requantize_node.meta = consumer.meta.copy()
+            requantize_node.meta["val"] = exir_ops.edge.cadence.requantize.per_tensor(
+                cast(torch.fx.Node, consumer.args[0]).meta["val"],
+                cast(float, in_scale),
+                cast(int, in_zero_point),
+                cast(float, out_scale),
+                cast(int, out_zero_point),
+                cast(torch.dtype, out_dtype),
+            )
+            requantize_node.meta["tensor_meta"] = None
         return requantize_node
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
@@ -737,7 +747,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             ),
         )
         if modified:
-            return super().call(graph_module)
+            graph_module.graph.eliminate_dead_code()
+            graph_module.recompile()
+            return PassResult(graph_module, True)
         return PassResult(graph_module, False)
 
 

backends/cadence/aot/remove_ops.py

@@ -613,7 +613,7 @@ def __init__(self) -> None:
         )
 
 
-class RemoveSqueezeViewBeforeElementwiseOps(ExportPass):
+class RemoveSqueezeViewBeforeElementwiseOps(PassBase):
     """
     Looks for subgraphs of the form:
     squeeze -> [elementwise ops] -> view
@@ -667,6 +667,17 @@ def get_squeeze_indices(self, view_node: Node) -> List[int]:
 
         return squeeze_indices
 
+    def recompute_meta(self, nodes: List[Node]) -> None:
+        for node in nodes:
+            args, kwargs = pytree.tree_map_only(
+                Node,
+                lambda arg: arg.meta["val"],
+                (node.args, node.kwargs),
+            )
+            assert callable(node.target)
+            node.meta["val"] = node.target(*args, **kwargs)
+            node.meta["tensor_meta"] = None
+
     def handle_squeeze(self, view_node: Node, visited_view_nodes: Set[Node]) -> bool:
         if view_node in visited_view_nodes:
             return False
@@ -681,13 +692,15 @@ def handle_squeeze(self, view_node: Node, visited_view_nodes: Set[Node]) -> bool
         node = next(iter(view_node.users))
 
         # Traverse down from the node until finding another view op.
+        intermediate_nodes = []
         intermediate_slices = []
         while node.target != exir_ops.edge.aten.view_copy.default:
             # Only handle simple chains for now
             if len(node.users) != 1:
                 return False
             if node.target not in self.intermediate_ops:
                 return False
+            intermediate_nodes.append(node)
             if node.target == exir_ops.edge.aten.slice_copy.Tensor:
                 intermediate_slices.append(node)
             node = next(iter(node.users))
@@ -710,6 +723,7 @@ def handle_squeeze(self, view_node: Node, visited_view_nodes: Set[Node]) -> bool
         # Skip the initial view node.
         input_node = get_arg(view_node, "input", Node)
         view_node.replace_all_uses_with(input_node)
+        self.recompute_meta(intermediate_nodes)
         return True
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
@@ -723,7 +737,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         if modified:
             graph_module.graph.eliminate_dead_code()
             graph_module.recompile()
-            return super().call(graph_module)
+            return PassResult(graph_module, True)
 
         return PassResult(graph_module, False)
 

backends/cadence/aot/tests/test_fusion_ops_passes.py

@@ -534,12 +534,28 @@ def test_replace_dequant_quant_with_requantize(self) -> None:
         )
         builder.output([quant])
         original_graph = builder.get_graph_module()
+        original_fx_graph = original_graph.graph
+        quant_node = original_graph.graph.find_nodes(
+            op="call_function",
+            target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+        )[0]
+        quant_node.meta["tier_2a_sentinel"] = object()
+        quant_sentinel = quant_node.meta["tier_2a_sentinel"]
         gm_before = copy.deepcopy(original_graph)
 
         p = FuseQuantDequantToRequantizePass()
         result = cast(PassResult, p(original_graph))
         self.assertTrue(result.modified)
+        self.assertIs(result.graph_module, original_graph)
+        self.assertIs(result.graph_module.graph, original_fx_graph)
         converted_graph = result.graph_module
+        requantize_node = converted_graph.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.cadence.requantize.per_tensor
+        )[0]
+        self.assertIs(requantize_node.meta["tier_2a_sentinel"], quant_sentinel)
+        self.assertIsNot(requantize_node.meta, quant_node.meta)
+        self.assertEqual(requantize_node.meta["val"].shape, x_input.shape)
+        self.assertEqual(requantize_node.meta["val"].dtype, torch.int8)
 
         # Validate numerical accuracy
         validate_numerics(
@@ -1201,12 +1217,27 @@ def test_fuse_transpose_permute_pairs(
         )
 
         # Check that the pass fuses the two transpose/permute ops.
+        original_fx_graph = gm.graph
+        output_node = next(
+            node for node in reversed(gm.graph.nodes) if node.op == "call_function"
+        )
+        output_node.meta["tier_2a_sentinel"] = object()
+        output_sentinel = output_node.meta["tier_2a_sentinel"]
+        output_shape = output_node.meta["val"].shape
         fusion_pass_result = FuseTransposeOrPermuteOpPairsPass()(gm)
         self.assertIsNotNone(fusion_pass_result)
+        self.assertIs(fusion_pass_result.graph_module, gm)
+        self.assertIs(fusion_pass_result.graph_module.graph, original_fx_graph)
         gm_after_pass = fusion_pass_result.graph_module
         if expected_is_fused:
             expected_op_counts[op1] = 0
             expected_op_counts[op2] = 0
+            replacement_view = gm_after_pass.graph.find_nodes(
+                op="call_function", target=exir_ops.edge.aten.view_copy.default
+            )[0]
+            self.assertIs(replacement_view.meta["tier_2a_sentinel"], output_sentinel)
+            self.assertIsNot(replacement_view.meta, output_node.meta)
+            self.assertEqual(replacement_view.meta["val"].shape, output_shape)
         self.check_op_counts(
             gm_after_pass,
             # pyre-fixme[6]: Incompatible parameter type

backends/cadence/aot/tests/test_remove_ops_passes.py

@@ -599,11 +599,14 @@ def test_remove_squeeze_view_before_elemwise_ops(self) -> None:
         )
         builder.output([unsqueeze])
         model = builder.get_graph_module()
+        original_graph = model.graph
         original = copy.deepcopy(model)
 
         p = RemoveSqueezeViewBeforeElementwiseOps()
         pass_result = cast(PassResult, p(model))
         self.assertTrue(pass_result.modified)
+        self.assertIs(pass_result.graph_module, model)
+        self.assertIs(pass_result.graph_module.graph, original_graph)
         transformed = pass_result.graph_module
 
         # First view should be eliminated and second view should be trivial.
@@ -619,6 +622,13 @@ def test_remove_squeeze_view_before_elemwise_ops(self) -> None:
         )
         self.assertEqual(len(slices), 1)
         self.assertEqual(slices[0].args[1], 2)
+        quantize_nodes = transformed.graph.find_nodes(
+            op="call_function",
+            target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+        )
+        self.assertEqual(quantize_nodes[0].meta["val"].shape, torch.Size([8, 1, 4, 4]))
+        self.assertEqual(slices[0].meta["val"].shape, torch.Size([8, 1, 2, 4]))
+        self.assertEqual(views[0].meta["val"].shape, torch.Size([8, 1, 2, 4]))
 
         # Verify the output of the model is the same as the original.
         sample_input = torch.randn(8, 1, 4, 4)
@@ -650,10 +660,15 @@ def test_remove_squeeze_view_before_elemwise_ops_multiple_squeeze(self) -> None:
         )
         builder.output([view_copy])
         model = builder.get_graph_module()
+        original_graph = model.graph
         original = copy.deepcopy(model)
 
         p = RemoveSqueezeViewBeforeElementwiseOps()
-        transformed = cast(PassResult, p(model)).graph_module
+        pass_result = cast(PassResult, p(model))
+        self.assertTrue(pass_result.modified)
+        self.assertIs(pass_result.graph_module, model)
+        self.assertIs(pass_result.graph_module.graph, original_graph)
+        transformed = pass_result.graph_module
 
         # First view should be eliminated.
         self.assertEqual(
@@ -666,6 +681,14 @@ def test_remove_squeeze_view_before_elemwise_ops_multiple_squeeze(self) -> None:
         )
         self.assertEqual(len(slices), 1)
         self.assertEqual(slices[0].args[1], 3)
+        quantize_nodes = transformed.graph.find_nodes(
+            op="call_function",
+            target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+        )
+        self.assertEqual(
+            quantize_nodes[0].meta["val"].shape, torch.Size([8, 1, 1, 4, 1, 4])
+        )
+        self.assertEqual(slices[0].meta["val"].shape, torch.Size([8, 1, 1, 2, 1, 4]))
 
         # Verify the output of the model is the same as the original.
         sample_input = torch.randn(8, 1, 1, 4, 1, 4)

backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py

@@ -105,5 +105,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             bypass_ops=self.bypass_ops,
         )
         if modified:
-            return super().call(graph_module)
+            graph_module.graph.eliminate_dead_code()
+            graph_module.recompile()
+            return PassResult(graph_module, True)
         return PassResult(graph_module, False)