diff --git a/backends/cadence/aot/fuse_ops.py b/backends/cadence/aot/fuse_ops.py index 51324f64630..26d42db6546 100644 --- a/backends/cadence/aot/fuse_ops.py +++ b/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) diff --git a/backends/cadence/aot/remove_ops.py b/backends/cadence/aot/remove_ops.py index 7603db75bf1..0ca01cb3133 100644 --- a/backends/cadence/aot/remove_ops.py +++ b/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,6 +692,7 @@ 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 @@ -688,6 +700,7 @@ def handle_squeeze(self, view_node: Node, visited_view_nodes: Set[Node]) -> bool 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) diff --git a/backends/cadence/aot/tests/test_fusion_ops_passes.py b/backends/cadence/aot/tests/test_fusion_ops_passes.py index 48ed25be423..646b7ad9fe6 100644 --- a/backends/cadence/aot/tests/test_fusion_ops_passes.py +++ b/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 diff --git a/backends/cadence/aot/tests/test_remove_ops_passes.py b/backends/cadence/aot/tests/test_remove_ops_passes.py index 1bce74d9ab9..afd777719b7 100644 --- a/backends/cadence/aot/tests/test_remove_ops_passes.py +++ b/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) diff --git a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py index 83b33533995..b8ae98d9335 100644 --- a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py +++ b/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)