diff --git a/backends/arm/_passes/rewrite_conv_pass.py b/backends/arm/_passes/rewrite_conv_pass.py index dc00baebaac..9bd6047f678 100644 --- a/backends/arm/_passes/rewrite_conv_pass.py +++ b/backends/arm/_passes/rewrite_conv_pass.py @@ -37,11 +37,20 @@ TOSA_CONTROL_FLOW_SOURCE_NODE_META, TosaSpecialDtype, ) -from executorch.backends.arm.tosa.specification import get_context_shape_env +from executorch.backends.arm.tosa.specification import ( + get_context_shape_env, + get_context_spec, +) +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + build_node_signature, + DO_NOT_FUSE_DUPLICATE_META_KEY, +) from executorch.backends.transforms.utils import create_constant_placeholder from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass, PassResult +from torch._ops import OpOverload from torch._subclasses.fake_tensor import FakeTensor from torch.export.graph_signature import InputKind @@ -533,7 +542,7 @@ def _is_direct_int32_rescale(node: torch.fx.Node) -> bool: def _get_direct_int32_rescale_users( self, node: torch.fx.Node - ) -> list[torch.fx.Node]: + ) -> list[torch.fx.Node] | None: """Return consumers that directly request an INT32 value.""" return [user for user in node.users if self._is_direct_int32_rescale(user)] @@ -571,6 +580,87 @@ def _insert_layout_permute( output.meta["val"] = output_fake_tensor return output, output_fake_tensor + @classmethod + def _deduplicate_a16w8_output_rescales( + cls, + graph_module: torch.fx.GraphModule, + tosa_op: torch.fx.Node, + node_order: dict[torch.fx.Node, int], + ) -> list[torch.fx.Node]: + """Merge only complete, canonical RESCALE-to-PERMUTE heads.""" + rescale_users = sorted( + tosa_op.users, key=lambda node: node_order.get(node, len(node_order)) + ) + if any( + user.target != exir_ops.backend.tosa.RESCALE.default + for user in rescale_users + ): + return None + + unique_rescales: dict[tuple[Any, ...], torch.fx.Node] = {} + deduplicated_rescales: list[torch.fx.Node] = [] + for rescale in rescale_users: + rescale_outputs = list(rescale.users) + if ( + len(rescale_outputs) != 1 + or rescale_outputs[0].target != exir_ops.edge.aten.permute_copy.default + ): + deduplicated_rescales.append(rescale) + continue + layout_permute = rescale_outputs[0] + rescale_signature = build_node_signature( + rescale, positional_arg_start=1 + ) + permute_signature = build_node_signature( + layout_permute, positional_arg_start=1 + ) + if rescale_signature is None or permute_signature is None: + deduplicated_rescales.append(rescale) + continue + signature = ( + rescale_signature, + permute_signature, + ) + canonical_permute = unique_rescales.get(signature) + if canonical_permute is not None: + # Layout permutes are inserted directly after their RESCALE, + # so the earliest RESCALE also provides a dominating permute. + layout_permute.replace_all_uses_with(canonical_permute) + graph_module.graph.erase_node(layout_permute) + graph_module.graph.erase_node(rescale) + else: + unique_rescales[signature] = layout_permute + deduplicated_rescales.append(rescale) + + return deduplicated_rescales + + def _separate_u55_a16w8_output_rescales( + self, + graph_module: torch.fx.GraphModule, + tosa_op: torch.fx.Node, + node_order: dict[torch.fx.Node, int], + ) -> None: + if len(tosa_op.users) < 2: + return + + rescale_users = self._deduplicate_a16w8_output_rescales( + graph_module, tosa_op, node_order + ) + if rescale_users is None or len(rescale_users) < 2: + return + tosa_op.meta[DO_NOT_FUSE_DUPLICATE_META_KEY] = True + for rescale in rescale_users[1:]: + with graph_module.graph.inserting_before(rescale): + cloned_tosa_op = create_node( + graph=graph_module.graph, + op_target=cast(OpOverload | EdgeOpOverload, tosa_op.target), + args=tosa_op.args, + kwargs=tosa_op.kwargs, + from_node=tosa_op, + inherit_qparams=True, + ) + rescale.replace_input_with(tosa_op, cloned_tosa_op) + def _insert_a16w8_output_branches( self, graph_module: torch.fx.GraphModule, @@ -787,6 +877,7 @@ def _insert_output_conversion( def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 modified = False + a16w8_tosa_ops: list[torch.fx.Node] = [] for node in graph_module.graph.nodes: if ( node.op != "call_function" @@ -1172,11 +1263,21 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 if squeeze_view is not None: graph_module.graph.erase_node(squeeze_view) graph_module.graph.erase_node(output_conversion_node) + a16w8_tosa_ops.append(tosa_op) else: node.replace_all_uses_with(node_replacement) graph_module.graph.erase_node(node) + if a16w8_tosa_ops and get_context_spec().is_U55_subset: + node_order = { + node: index for index, node in enumerate(graph_module.graph.nodes) + } + for tosa_op in a16w8_tosa_ops: + self._separate_u55_a16w8_output_rescales( + graph_module, tosa_op, node_order + ) + if modified: graph_module.recompile() graph_module = super().call(graph_module).graph_module diff --git a/backends/arm/test/passes/test_fuse_duplicate_users_pass.py b/backends/arm/test/passes/test_fuse_duplicate_users_pass.py index 893d9eefea5..831d9d91267 100644 --- a/backends/arm/test/passes/test_fuse_duplicate_users_pass.py +++ b/backends/arm/test/passes/test_fuse_duplicate_users_pass.py @@ -21,6 +21,9 @@ TosaLoweringContext, TosaSpecification, ) +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + DO_NOT_FUSE_DUPLICATE_META_KEY, +) from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export @@ -163,6 +166,22 @@ def test_fuse_duplicate_users_preserves_graph_order_for_representative(): assert len(_add_node_names(result.graph_module)) == 1 +def test_fuse_duplicate_users_honors_do_not_fuse_marker(): + graph_module = _graph_with_users_not_in_node_order() + marked_node = next( + node + for node in graph_module.graph.nodes + if node.target == torch.ops.aten.add.Tensor + ) + marked_node.meta[DO_NOT_FUSE_DUPLICATE_META_KEY] = True + + result = FuseDuplicateUsersPass()(graph_module) + + result.graph_module.graph.lint() + assert not result.modified + assert len(_add_node_names(result.graph_module)) == 2 + + def test_fuse_duplicate_users_keeps_identical_rescale_users(): graph_module = _graph_with_duplicate_rescale_users() diff --git a/backends/arm/test/passes/test_rewrite_conv_pass.py b/backends/arm/test/passes/test_rewrite_conv_pass.py index e1605faefa6..60ccd412aa5 100644 --- a/backends/arm/test/passes/test_rewrite_conv_pass.py +++ b/backends/arm/test/passes/test_rewrite_conv_pass.py @@ -27,7 +27,10 @@ from executorch.backends.arm.test.misc.test_dw_convs_with_shared_weights import ( DWConvsModule, ) -from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, + PassPipeline, +) from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec from executorch.backends.arm.tosa.mapping import TosaSpecialDtype from executorch.backends.arm.tosa.partitioner import TOSAPartitioner @@ -36,6 +39,9 @@ TosaSpecification, ) from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + build_node_signature, +) from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower from executorch.exir.dialects._ops import ops as exir_ops from torch.export import Dim, export @@ -267,7 +273,9 @@ def _get_expected_int32_scales( def _rewrite_a16w8_convs( - model: nn.Module, inputs: tuple[torch.Tensor, ...] + model: nn.Module, + inputs: tuple[torch.Tensor, ...], + tosa_spec: TosaSpecification | None = None, ) -> tuple[torch.fx.GraphModule, list[list[float]]]: """Run the passes needed to inspect rewritten A16W8 convolutions.""" exported_program = _export_quantized_a16w8(model, inputs) @@ -276,7 +284,7 @@ def _rewrite_a16w8_convs( ).exported_program() gm = _run_pre_rewrite_passes(edge_program) rewrite_pass = RewriteConvPass(edge_program) - with TosaLoweringContext(_compile_spec_int16().tosa_spec): + with TosaLoweringContext(tosa_spec or _compile_spec_int16().tosa_spec): rescale_result = InsertRescaleInt32Pass()(gm) assert rescale_result is not None expected_int32_scales = _get_expected_int32_scales( @@ -294,6 +302,29 @@ def _get_call_function_node(gm: torch.fx.GraphModule, target): raise AssertionError(f"Node with target {target} not found") +def _add_a16w8_rescale_head( + graph: torch.fx.Graph, + accumulator: torch.fx.Node, + positional_unsigned: tuple[bool, ...] = (), +) -> tuple[torch.fx.Node, torch.fx.Node]: + rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + args=( + accumulator, + torch.int16, + [1.0], + 0, + 0, + *positional_unsigned, + ), + ) + layout_permute = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(rescale, [0, 3, 1, 2]), + ) + return rescale, layout_permute + + class ConvModule(torch.nn.Module): def __init__(self): super().__init__() @@ -510,6 +541,149 @@ def test_rewrite_conv_a16w8_mixed_consumers_restore_int16( ) +def test_rewrite_conv_rescale_signature_includes_positional_unsigned_flags() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + signed_rescale, _ = _add_a16w8_rescale_head(graph, accumulator, (False, False)) + unsigned_rescale, _ = _add_a16w8_rescale_head(graph, accumulator, (False, True)) + + assert build_node_signature( + signed_rescale, positional_arg_start=1 + ) != build_node_signature(unsigned_rescale, positional_arg_start=1) + + +def test_rewrite_conv_without_convolution_does_not_require_context() -> None: + inputs = (torch.randn(1, 4),) + edge_program = to_edge(export(nn.Identity(), inputs)).exported_program() + + result = RewriteConvPass(edge_program)(edge_program.graph_module) + + assert result is not None + assert not result.modified + + +def test_rewrite_conv_a16w8_unknown_accumulator_user_is_unchanged() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + _, first_permute = _add_a16w8_rescale_head(graph, accumulator) + _, second_permute = _add_a16w8_rescale_head(graph, accumulator) + unexpected_user = graph.call_function(torch.neg, args=(accumulator,)) + graph.output((first_permute, second_permute, unexpected_user)) + graph_module = torch.fx.GraphModule({}, graph) + nodes_before = list(graph.nodes) + users_before = {node: tuple(node.users) for node in graph.nodes} + node_order = {node: index for index, node in enumerate(graph.nodes)} + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result is None + assert list(graph.nodes) == nodes_before + assert {node: tuple(node.users) for node in graph.nodes} == users_before + graph.lint() + + +def test_rewrite_conv_a16w8_multi_consumer_rescale_is_not_deduplicated() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + first_rescale, first_permute = _add_a16w8_rescale_head(graph, accumulator) + second_rescale, second_permute = _add_a16w8_rescale_head(graph, accumulator) + output = graph.output((first_permute, second_permute, second_rescale)) + graph_module = torch.fx.GraphModule({}, graph) + node_order = {node: index for index, node in enumerate(graph.nodes)} + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result == [first_rescale, second_rescale] + assert set(second_rescale.users) == {second_permute, output} + graph.lint() + + +def test_rewrite_conv_a16w8_deduplication_uses_graph_order() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + temporary_input = graph.placeholder("temporary_input") + early_rescale, early_permute = _add_a16w8_rescale_head(graph, temporary_input) + late_rescale, late_permute = _add_a16w8_rescale_head(graph, accumulator) + early_rescale.replace_input_with(temporary_input, accumulator) + graph.output((early_permute, late_permute)) + graph_module = torch.fx.GraphModule({}, graph) + node_order = {node: index for index, node in enumerate(graph.nodes)} + + assert list(accumulator.users) == [late_rescale, early_rescale] + assert node_order[early_rescale] < node_order[late_rescale] + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result == [early_rescale] + assert late_rescale not in graph.nodes + assert late_permute not in graph.nodes + graph.lint() + + +def test_rewrite_conv_a16w8_deduplication_tolerates_new_users() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + first_rescale, first_permute = _add_a16w8_rescale_head(graph, accumulator) + node_order = {node: index for index, node in enumerate(graph.nodes)} + second_rescale, second_permute = _add_a16w8_rescale_head(graph, accumulator) + graph.output((first_permute, second_permute)) + graph_module = torch.fx.GraphModule({}, graph) + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result == [first_rescale] + assert second_rescale not in graph.nodes + graph.lint() + + +def test_rewrite_conv_a16w8_u55_separates_distinct_output_rescales() -> None: + model = A16W8MixedConsumerChain(nn.Conv2d(4, 4, 1)) + inputs = (torch.randn(1, 4, 8, 8),) + generic_graph, _ = _rewrite_a16w8_convs(model, inputs) + u55_graph, _ = _rewrite_a16w8_convs( + model, + inputs, + TosaSpecification.create_from_string("TOSA-1.0+INT+int16+int4+u55"), + ) + + conv_targets = { + exir_ops.backend.tosa.CONV2D.default, + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default, + } + generic_convs = [ + node for node in generic_graph.graph.nodes if node.target in conv_targets + ] + u55_convs = [node for node in u55_graph.graph.nodes if node.target in conv_targets] + + assert len(u55_convs) == len(generic_convs) + 1 + assert all(len(conv.users) == 1 for conv in u55_convs) + assert all( + next(iter(conv.users)).target == exir_ops.backend.tosa.RESCALE.default + for conv in u55_convs + ) + + +def test_rewrite_conv_a16w8_mixed_consumers_lowers_on_u55() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + pipeline = EthosU55PipelineINT[tuple[torch.Tensor]]( + A16W8MixedConsumerChain(nn.Conv2d(4, 4, 1)), + inputs, + aten_ops=[], + exir_ops=[], + run_on_fvp=False, + a16w8_quantization=True, + ) + pipeline.run() + + def test_rewrite_conv_a16w8_preserves_int32_for_int32_consumers() -> None: r"""Test that an exclusively INT32 consumer keeps the widened path. diff --git a/backends/transforms/fuse_duplicate_users_pass.py b/backends/transforms/fuse_duplicate_users_pass.py index b3989e76c94..47b6064b618 100644 --- a/backends/transforms/fuse_duplicate_users_pass.py +++ b/backends/transforms/fuse_duplicate_users_pass.py @@ -11,7 +11,76 @@ from executorch.exir.pass_base import ExportPass, PassResult from torch._ops import OpOverload from torch.fx import GraphModule, Node -from torch.fx.node import Argument, map_arg +from torch.fx.node import map_arg + + +DO_NOT_FUSE_DUPLICATE_META_KEY = "do_not_fuse_duplicate" + + +def _map_leaf_to_key(node: Node) -> str: + return node.name + + +def _to_hashable(value: Any) -> Hashable: + """Convert arbitrarily nested structures into hashable tuples.""" + + if isinstance(value, (list, tuple)): + return tuple(_to_hashable(v) for v in value) + if isinstance(value, dict): + normalized_items = [(k, _to_hashable(v)) for k, v in value.items()] + return tuple(sorted(normalized_items, key=lambda item: repr(item[0]))) + if isinstance(value, set): + hashable_values: List[Hashable] = [_to_hashable(v) for v in value] + return tuple(sorted(hashable_values, key=repr)) + if isinstance(value, slice): + return ( + "slice", + _to_hashable(value.start), + _to_hashable(value.stop), + _to_hashable(value.step), + ) + if isinstance(value, range): + return ("range", value.start, value.stop, value.step) + if isinstance(value, torch.Size): + return ("size", tuple(value)) + if isinstance(value, torch.dtype): + return ("dtype", str(value)) + if isinstance(value, torch.device): + return ("device", str(value)) + if isinstance(value, torch.memory_format): + return ("memory_format", str(value)) + if isinstance(value, torch.Tensor): + return ( + "tensor", + str(value.dtype), + tuple(value.size()), + value.device.type, + value.requires_grad, + ) + return value + + +def _get_target_key(target: Any) -> Hashable: + if isinstance(target, (EdgeOpOverload, OpOverload)): + return str(target) + return target + + +def build_node_signature( + node: Node, *, positional_arg_start: int = 0 +) -> Tuple[Hashable, ...] | None: + """Build a stable signature while ignoring leading positional operands.""" + try: + normalized_args = _to_hashable( + map_arg(node.args[positional_arg_start:], _map_leaf_to_key) + ) + normalized_kwargs = _to_hashable( + {k: map_arg(v, _map_leaf_to_key) for k, v in node.kwargs.items()} + ) + except TypeError: + return None + + return (node.op, _get_target_key(node.target), normalized_args, normalized_kwargs) class FuseDuplicateUsersPass(ExportPass): @@ -106,7 +175,7 @@ def _get_candidate_groups(self, node_order, user_nodes): if user.target in self._excluded_targets: continue - target_key = self._get_target_key(user.target) + target_key = _get_target_key(user.target) target_signature = (user.op, target_key) users_by_target.setdefault(target_signature, []).append(user) @@ -120,62 +189,6 @@ def _get_candidate_groups(self, node_order, user_nodes): return candidate_groups def _build_user_signature(self, node: Node) -> Tuple[Hashable, ...] | None: - try: - normalized_args = self._to_hashable( - map_arg(node.args, self._map_leaf_to_key) - ) - normalized_kwargs = self._to_hashable( - {k: map_arg(v, self._map_leaf_to_key) for k, v in node.kwargs.items()} - ) - except TypeError: + if node.meta.get(DO_NOT_FUSE_DUPLICATE_META_KEY, False): return None - - target_key = self._get_target_key(node.target) - - return (node.op, target_key, normalized_args, normalized_kwargs) - - def _map_leaf_to_key(self, node: Node) -> Argument: - return node.name - - def _to_hashable(self, value: Any) -> Hashable: - """Convert arbitrarily nested structures into hashable tuples.""" - - if isinstance(value, (list, tuple)): - return tuple(self._to_hashable(v) for v in value) - if isinstance(value, dict): - normalized_items = [(k, self._to_hashable(v)) for k, v in value.items()] - return tuple(sorted(normalized_items, key=lambda item: repr(item[0]))) - if isinstance(value, set): - hashable_values: List[Hashable] = [self._to_hashable(v) for v in value] - return tuple(sorted(hashable_values, key=repr)) - if isinstance(value, slice): - return ( - "slice", - self._to_hashable(value.start), - self._to_hashable(value.stop), - self._to_hashable(value.step), - ) - if isinstance(value, range): - return ("range", value.start, value.stop, value.step) - if isinstance(value, torch.Size): - return ("size", tuple(value)) - if isinstance(value, torch.dtype): - return ("dtype", str(value)) - if isinstance(value, torch.device): - return ("device", str(value)) - if isinstance(value, torch.memory_format): - return ("memory_format", str(value)) - if isinstance(value, torch.Tensor): - return ( - "tensor", - str(value.dtype), - tuple(value.size()), - value.device.type, - value.requires_grad, - ) - return value - - def _get_target_key(self, target: Any) -> Hashable: - if isinstance(target, (EdgeOpOverload, OpOverload)): - return str(target) - return target + return build_node_signature(node)