[Group Partitioner] leverage group partitioner for config-based partitioner

ghstack-source-id: fdc99b6
ghstack-comment-id: 3115642963
Pull Request resolved: #12845

Author: mcr229

exir/backend/canonical_partitioners/config_partitioner.py

@@ -10,13 +10,17 @@
 import torch
 from executorch.exir.backend.backend_details import ExportedProgram
 from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import (
-    generate_partitions_from_list_of_nodes,
+    generate_grouped_partitions_from_list_of_nodes,
 )
 from executorch.exir.backend.partitioner import (
     DelegationSpec,
     Partitioner,
     PartitionResult,
 )
+
+from exir.backend.canonical_partitioners.pattern_op_partitioner import (
+    generate_grouped_partitions_from_list_of_nodes,
+)
 from torch.fx.passes.infra.partitioner import Partition
 
 
@@ -162,23 +166,49 @@ def filter_fn(node: torch.fx.Node) -> bool:
     def get_matched_nodes_from_configs(
         self, ep: ExportedProgram
     ) -> List[List[torch.fx.Node]]:
+        # disjoint set union
+        parent = {}
+
+        def find(x):
+            parent.setdefault(x, x)
+            if parent[x] != x:
+                parent[x] = find(parent[x])
+            return parent[x]
+
+        def union(x, y):
+            parent[find(x)] = find(y)
+
         # gather supported nodes
-        matched_nodes = []
         gm = ep.graph_module
         for node in gm.graph.nodes:
-            if node.op == "call_function":
-                target = format_target_name(node.target.__name__)
-                if target in self.target_partitioner_configs:
-                    node_config = self.target_partitioner_configs[target]
-                    if node_config.check_constraints(node, ep):
-                        matched_nodes.append(node_config.get_partition(node, ep))
+            if node.op != "call_function":
+                continue
+            target = format_target_name(node.target.__name__)
+
+            if target not in self.target_partitioner_configs:
+                continue
+
+            node_config = self.target_partitioner_configs[target]
+            if not node_config.check_constraints(node, ep):
+                continue
+
+            partition = node_config.get_partition(node, ep)
+            if len(partition) > 0:
+                parent[partition[0]] = partition[0]
+                for i in range(1, len(partition)):
+                    union(partition[0], partition[i])
+
+        groups = {}
+        for node in parent.keys():
+            root = find(node)
+            groups.setdefault(root, set()).add(node)
 
-        return matched_nodes
+        return [list(group) for group in groups.values()]
 
     def generate_partitions(self, ep: ExportedProgram) -> List[Partition]:
         matched_nodes = self.get_matched_nodes_from_configs(ep)
         # create partitions
-        partitions = generate_partitions_from_list_of_nodes(
+        partitions = generate_grouped_partitions_from_list_of_nodes(
             ep.graph_module,
             matched_nodes,
         )

exir/backend/canonical_partitioners/pattern_op_partitioner.py

@@ -8,6 +8,10 @@
 from typing import List, Optional
 
 import torch
+
+from executorch.exir.backend.canonical_partitioners.group_partitioner import (
+    GroupBasedPartitioner,
+)
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import any_chain, OperatorSupportBase
 from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
@@ -56,6 +60,50 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
     return partition_list
 
 
+def generate_grouped_partitions_from_list_of_nodes(
+    graph_module: torch.fx.GraphModule,
+    pattern_list: Optional[List[List[torch.fx.Node]]] = None,
+    op_support: Optional[OperatorSupportBase] = None,
+) -> List[Partition]:
+    final_op_support: Optional[OperatorSupportBase] = op_support
+
+    if pattern_list is not None:
+        # Tag all the nodes in these patterns
+        for node_list in pattern_list:
+            for node in node_list:
+                node.meta["match"] = True
+
+        class MatchTag(OperatorSupportBase):
+            def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
+                return node.meta.get("match", False)
+
+        final_op_support = (
+            MatchTag()
+            if final_op_support is None
+            else any_chain(final_op_support, MatchTag())
+        )
+
+    assert (
+        final_op_support is not None
+    ), "Did not give a pattern or OperatorSupportBase instance to partition with"
+
+    # Run the CapabilityBasedPartitioner to return the largest possible
+    # subgraphs containing the nodes with the tags
+    group_partitioner = GroupBasedPartitioner(
+        graph_module,
+        final_op_support,
+        node_groups=pattern_list,
+        allows_single_node_partition=True,
+    )
+    partition_list = group_partitioner.propose_partitions()
+
+    # Remove the metadata field we added
+    for partition in partition_list:
+        for node in partition.nodes:
+            node.meta.pop("match", False)
+    return partition_list
+
+
 def generate_pattern_op_partitions(
     graph_module: torch.fx.GraphModule,
     patterns: Optional[List[torch.fx.Graph]] = None,