Arm backend: Tag control flow submodules in partitioner (#15364)

Control flow operators are implemented with get_attr/ placeholders nodes
pointing to submodules in the graph_module. To know if we can handle
such an operators, we need to figure out whether we can partition all of
the submodule the operator points to. We can do this simply by applying
the partition logic we already have to submodules. This has the added
benefit that we will partition parts of submodules into normal
partitions, even if we can't partition the full submodule.

Signed-off-by: Erik Lundell <erik.lundell@arm.com>

Author: Erik-Lundell

backends/arm/tosa/partitioner.py

@@ -15,6 +15,7 @@
 """
 
 import logging
+from itertools import count
 from typing import Callable, List, Optional, Sequence, Tuple
 
 import torch
@@ -35,8 +36,10 @@
 )
 from executorch.exir.backend.utils import tag_constant_data, WhyNoPartitionReporter
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.graph_module import get_control_flow_submodules
 from torch.export.exported_program import ExportedProgram
-from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
+from torch.fx import GraphModule
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
 
 logger = logging.getLogger(__name__)
@@ -110,6 +113,43 @@ def is_noop_expand(node: torch.fx.node.Node) -> bool:
     return all(m == 1 for m in multiples)
 
 
+def is_partitioned(
+    node: torch.fx.Node,
+    tag: str,
+) -> bool:
+    """Return True if the node currently belongs to the partition ``tag``.
+
+    Args:
+        node (torch.fx.Node): FX node to check.
+        tag (str): Delegation tag identifying the partition.
+
+    Returns:
+        bool: True if the node carries the matching delegation tag.
+
+    """
+    return "delegation_tag" in node.meta and node.meta["delegation_tag"] == tag
+
+
+def reject_partition(
+    reason: str, partition: Partition, reporter: WhyNoPartitionReporter
+) -> None:
+    """Remove a proposed partition and record the rejection reason.
+
+    Args:
+        reason (str): Human-readable explanation for rejection.
+        partition (object): Proposed partition object from the
+            capability partitioner.
+        reporter (WhyNoPartitionReporter): used to report why nodes were rejected.
+    """
+    for node in partition.nodes:
+        if "delegation_tag" in node.meta:
+            del node.meta["delegation_tag"]
+            reporter.report_reject(
+                node,
+                reason,
+            )
+
+
 class TOSAPartitioner(Partitioner):
     """Partition an exported program into TOSA-delegable subgraphs.
 
@@ -142,107 +182,76 @@ def __init__(
         self.additional_checks = additional_checks
         self.tosa_spec = compile_spec.tosa_spec
 
-    def partition(self, exported_program: ExportedProgram) -> PartitionResult:  # noqa
-        """Partition the program and tag TOSA-compatible subgraphs.
-
-        Run the FX capability-based partitioner to propose subgraphs, then
-        refine tags by removing boundary-only quantize/dequantize nodes and by
-        rejecting partitions that would lower to no-ops. Emit a detailed report
-        of rejected nodes and their reasons.
+    def _tag_module(  # noqa
+        self,
+        module: GraphModule,
+        containing_program: ExportedProgram,
+        reporter: WhyNoPartitionReporter,
+        tag_iterator: count | None = None,
+    ) -> set[str]:
+        """Tag nodes in a module, possibly a submodule, from the containing program.
 
         Args:
-            exported_program (ExportedProgram): Program to analyze and
-                partition.
-
+            module: a GraphModule from `containing_program` to tag nodes in.
+            containing_program: The ExportedProgram that contains the module.
+            reporter: A reporter to report why nodes were rejected.
         Returns:
-            PartitionResult: The input program with nodes tagged for delegation
-            and a mapping of partition tags to delegation specs.
-
+            A set of strings with the partition tags.
         """
-        logger.info("TOSAPartitioner::partition")
-        partition_tags: dict[str, DelegationSpec] = {}
-
-        logger.info(
-            f"Partitioning for {self.delegation_spec.backend_id}: {self.tosa_spec}"
-        )
-
-        reporter = WhyNoPartitionReporter()
+        tags: set[str] = set()
+        if tag_iterator is None:
+            tag_iterator = count(0)
+        for _, submodule, _ in get_control_flow_submodules(module):
+            submodule_tags = self._tag_module(
+                submodule, containing_program, reporter, tag_iterator
+            )
+            if len(tags & submodule_tags) != 0:
+                raise RuntimeError(
+                    "Got overlapping tags in two different modules, this shouldn't happen."
+                )
+            tags = tags | submodule_tags
         operator_support = tosa_support_factory(
-            self.tosa_spec, exported_program, reporter, self.additional_checks
+            self.tosa_spec, containing_program, reporter, self.additional_checks
         )
         capability_partitioner = CapabilityBasedPartitioner(
-            exported_program.graph_module,
+            module,
             operator_support,
             allows_single_node_partition=True,
         )
         partition_list = capability_partitioner.propose_partitions()
 
-        def reject_partition(reason: str, partition, tag) -> None:
-            """Remove a proposed partition and record the rejection reason.
-
-            Args:
-                reason (str): Human-readable explanation for rejection.
-                partition (object): Proposed partition object from the
-                    capability partitioner.
-                tag (str): Delegation tag associated with the partition.
-
-            """
-            for node in partition.nodes:
-                if "delegation_tag" in node.meta:
-                    del node.meta["delegation_tag"]
-                    reporter.report_reject(
-                        node,
-                        reason,
-                    )
-            partition_tags.pop(tag, None)
-
         for partition in partition_list:
-            tag = f"tag{partition.id}"
-
-            def is_partitioned(node: torch.fx.Node, tag=tag) -> bool:
-                """Return True if the node currently belongs to the partition ``tag``.
-
-                Args:
-                    node (torch.fx.Node): FX node to check.
-                    tag (str): Delegation tag identifying the partition.
-
-                Returns:
-                    bool: True if the node carries the matching delegation tag.
-
-                """
-                return (
-                    "delegation_tag" in node.meta and node.meta["delegation_tag"] == tag
-                )
+            tag = f"tag{next(tag_iterator)}"
+            tags.add(tag)
 
             for node in partition.nodes:
                 node.meta["delegation_tag"] = tag
-                partition_tags[tag] = self.delegation_spec
 
             # De-tag outermost q-nodes upwards and dq-nodes downwards.
             # De-tag if at least one input/output is not part of the partition.
-            for node in exported_program.graph_module.graph.nodes:
-                if not is_partitioned(node):
+            for node in module.graph.nodes:
+                if not is_partitioned(node, tag):
                     continue
                 if node.target in Q_OPS:
                     for input in node.all_input_nodes:
-                        if not is_partitioned(input):
+                        if not is_partitioned(input, tag):
                             del node.meta["delegation_tag"]
                             break
                     continue
 
                 if node.target in DQ_OPS:
                     for user in node.users:
-                        if not is_partitioned(user):
+                        if not is_partitioned(user, tag):
                             del node.meta["delegation_tag"]
                             break
                     continue
 
                 if self.tosa_spec.support_float():
                     continue
 
-                if is_partitioned(node):
+                if is_partitioned(node, tag):
                     for input in node.all_input_nodes:
-                        if is_partitioned(input):
+                        if is_partitioned(input, tag):
                             continue
                         if get_first_fake_tensor(input).dtype.is_floating_point:
                             reporter.report_reject(
@@ -265,8 +274,38 @@ def is_partitioned(node: torch.fx.Node, tag=tag) -> bool:
                 reject_partition(
                     "Partition contained only ops which are removed in the TOSA lowering, leading to an empty partition.",
                     partition,
-                    tag,
+                    reporter,
                 )
+                tags.remove(tag)
+        return tags
+
+    def partition(self, exported_program: ExportedProgram) -> PartitionResult:
+        """Partition the program and tag TOSA-compatible subgraphs.
+
+        Run the FX capability-based partitioner to propose subgraphs, then
+        refine tags by removing boundary-only quantize/dequantize nodes and by
+        rejecting partitions that would lower to no-ops. Emit a detailed report
+        of rejected nodes and their reasons.
+
+        Args:
+            exported_program (ExportedProgram): Program to analyze and
+                partition.
+
+        Returns:
+            PartitionResult: The input program with nodes tagged for delegation
+            and a mapping of partition tags to delegation specs.
+
+        """
+        logger.info("TOSAPartitioner::partition")
+        logger.info(
+            f"Partitioning for {self.delegation_spec.backend_id}: {self.tosa_spec}"
+        )
+
+        reporter = WhyNoPartitionReporter()
+        tags = self._tag_module(
+            exported_program.graph_module, exported_program, reporter
+        )
+        partition_tags = {tag: self.delegation_spec for tag in tags}
 
         tag_constant_data(exported_program)
         logger.info(f"The following nodes were rejected for {self.tosa_spec}:")