Wrap partition tags as part of partition result (#269)

Summary:
Pull Request resolved: #269

`partition_tags` is an output from the `partition` function and it'll be cleaner to wrap it as part of the outputs instead of the instance attributes. This diff is supposed be a noop

Reviewed By: mergennachin, angelayi

Differential Revision: D49116131

fbshipit-source-id: 0943c6367e2aa902e28af7f058bd2904f871beac

Author: cccclai

backends/qnnpack/partition/qnnpack_partitioner.py

@@ -19,7 +19,11 @@
 from executorch.backends.transforms.addmm_mm_to_linear import (
     apply_addmm_mm_to_linear_transform,
 )
-from executorch.exir.backend.partitioner import DelegationSpec, Partitioner
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
 from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
 
 logging.basicConfig(level=logging.INFO)
@@ -35,7 +39,6 @@ def __init__(self, delegate_name, patterns):
         self.patterns = patterns
 
         self.delegation_spec = DelegationSpec(delegate_name, [])
-        self.partition_tags: Dict[str, DelegationSpec] = {}
 
     @staticmethod
     def check_partitions(partitions: Union[dict, list]) -> None:
@@ -48,9 +51,9 @@ def check_partitions(partitions: Union[dict, list]) -> None:
         else:
             log.info(f"Found {pl} subgraphs to be partitioned.")
 
-    def partition(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    def partition(self, graph_module: torch.fx.GraphModule) -> PartitionResult:
         raise NotImplementedError("This is not meant to be used directly.")
-        return graph_module
+        return PartitionResult(tagged_graph=graph_module, partition_tags={})
 
 
 class _SingleOpDelegatePartitioner(_BasePartitioner):
@@ -72,7 +75,7 @@ def __init__(
         self.transforms = transforms
 
     # override
-    def partition(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    def partition(self, graph_module: torch.fx.GraphModule) -> PartitionResult:
         # TODO delete this since we are not allowed to do this
         if self.transforms is not None:
             for transform in self.transforms:  # pyre-ignore
@@ -107,7 +110,7 @@ def partition(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
 
         # Mapping from delegation tag to match set
         tag_mapping = {}
-
+        partition_tags: Dict[str, DelegationSpec] = {}
         for (partition_id, match_set) in enumerate(match_sets):
             delegation_tag = f"tag{partition_id}"
             for node in match_set:
@@ -124,10 +127,10 @@ def partition(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
                         )
                     break
                 node.meta["delegation_tag"] = delegation_tag
-            self.partition_tags[delegation_tag] = self.delegation_spec
+            partition_tags[delegation_tag] = self.delegation_spec
             tag_mapping[delegation_tag] = match_set
 
-        return graph_module
+        return PartitionResult(tagged_graph=graph_module, partition_tags=partition_tags)
 
 
 class QnnpackPartitioner(_SingleOpDelegatePartitioner):

backends/xnnpack/partition/xnnpack_partitioner.py

@@ -26,7 +26,11 @@
     generate_partitions_from_list_of_nodes,
 )
 
-from executorch.exir.backend.partitioner import DelegationSpec, Partitioner
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx.passes.infra.partitioner import Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
@@ -313,7 +317,6 @@ def __init__(
         self.supported_ops = set(supported_ops or [])
 
         self.delegation_spec = DelegationSpec(XnnpackBackend.__name__, [])
-        self.partition_tags: Dict[str, DelegationSpec] = {}
 
     @staticmethod
     def check_partitions(partitions: Union[dict, list]) -> bool:
@@ -386,27 +389,30 @@ def generate_partitions(self, graph_module: torch.fx.GraphModule) -> List[Any]:
             ),
         )
 
-    def tag_nodes(self, partitions: List[Partition]) -> None:
+    def tag_nodes(self, partitions: List[Partition]) -> Dict[str, DelegationSpec]:
         """
         Tag each partition in the list with its delegation tag.
         """
+        partition_tags: Dict[str, DelegationSpec] = {}
         for partition in partitions:
             # Add delegation tags
             for node in partition.nodes:
                 delegation_tag = f"tag{partition.id}"
                 node.meta["delegation_tag"] = delegation_tag
-                self.partition_tags[delegation_tag] = self.delegation_spec
+                partition_tags[delegation_tag] = self.delegation_spec
+        return partition_tags
 
     # override
-    def partition(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    def partition(self, graph_module: torch.fx.GraphModule) -> PartitionResult:
         """
         Run the partitioner on the given graph module, then tag each partition
         with its delegation tag (and partition id)
         """
         partitions = self.generate_partitions(graph_module)
+        partition_tags: Dict[str, DelegationSpec] = {}
         if self.check_partitions(partitions):
-            self.tag_nodes(partitions)
-        return graph_module
+            partition_tags = self.tag_nodes(partitions)
+        return PartitionResult(tagged_graph=graph_module, partition_tags=partition_tags)
 
 
 # TODO: Merge XnnpackQuantizedPartitioner and XnnpackFloatingPointPartitioner
@@ -761,10 +767,11 @@ def generate_partitions(
             XnnpackOperatorSupport(supported_ops=list(self.get_supported_ops(quant))),
         )
 
-    def tag_nodes(self, partitions: List[Partition]) -> None:
+    def tag_nodes(self, partitions: List[Partition]) -> Dict[str, DelegationSpec]:
         """
         Tag each partition in the list with its delegation tag.
         """
+        partition_tags: Dict[str, DelegationSpec] = {}
         for partition in partitions:
             # Add delegation tags
             skip = False
@@ -776,23 +783,25 @@ def tag_nodes(self, partitions: List[Partition]) -> None:
             for node in partition.nodes:
                 delegation_tag = f"tag{partition.id}"
                 node.meta["delegation_tag"] = delegation_tag
-                self.partition_tags[delegation_tag] = self.delegation_spec
+                partition_tags[delegation_tag] = self.delegation_spec
+        return partition_tags
 
     # override
     def _partition(
         self, graph_module: torch.fx.GraphModule, quant: Optional[bool]
-    ) -> torch.fx.GraphModule:
+    ) -> PartitionResult:
         """
         Run the partitioner on the given graph module, then tag each partition
         with its delegation tag (and partition id)
         """
         partitions = self.generate_partitions(graph_module, quant)
+        partition_tags: Dict[str, DelegationSpec] = {}
         if self.check_partitions(partitions):
-            self.tag_nodes(partitions)
-        return graph_module
+            partition_tags = self.tag_nodes(partitions)
+        return PartitionResult(tagged_graph=graph_module, partition_tags=partition_tags)
 
-    def partition(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
-        ret = self._partition(graph_module, self.quant)
+    def partition(self, graph_module: torch.fx.GraphModule) -> PartitionResult:
+        ret: PartitionResult = self._partition(graph_module, self.quant)
         return ret
 
 
@@ -805,7 +814,7 @@ def __init__(
         super().__init__(supported_modules, supported_ops)
 
     # override
-    def partition(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    def partition(self, graph_module: torch.fx.GraphModule) -> PartitionResult:
         """
         Run the partitioner on the given graph module, then tag each partition with its delegegation tag (and partition id)
 
@@ -819,7 +828,8 @@ def partition(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
             )
             for match in self.get_module_partitions(graph_module)
         ]
+        partition_tags: Dict[str, DelegationSpec] = {}
 
         if self.check_partitions(partitions):
-            self.tag_nodes(partitions)
-        return graph_module
+            partition_tags = self.tag_nodes(partitions)
+        return PartitionResult(tagged_graph=graph_module, partition_tags=partition_tags)

docs/website/docs/tutorials/backend_delegate.md

@@ -205,7 +205,7 @@ def to_backend(
 ```
 
 This function takes in a `Partitioner` which adds a tag to all the nodes that
-are meant to be lowered. It will also contain a `partition_tags` mapping tags to
+are meant to be lowered. The `Partitioner.partition` function will return both the tagged graph module and the `partition_tags` mapping tags to
 backend names and module compile specs. The tagged nodes will then be
 partitioned and lowered to their mapped backends using Flow 1's process.
 Available helper partitioner are documented [here](./passes.md#partitioner). These
@@ -371,24 +371,25 @@ class Backend_1_2_Partitioner(Partitioner):
     def __init__(self) -> None:
         self.delegation_spec_1 = DelegationSpec("Backend1", [])
         self.delegation_spec_2 = DelegationSpec("Backend2", [])
-        self.partition_tags = {}
 
     def partition(
         self, edge_graph_module: torch.fx.GraphModule
-    ) -> torch.fx.GraphModule:
-
+    ) -> PartitionResult:
+        partition_tags: Dict[str, DelegationSpec] = {}
         # Tag all nodes in the first partiton to backend 1
         node_to_backend_1 = ... # some logic to select the nodes from the graph
         delegation_tag = f"backend2_tag{partitioner_1.id}"
         node.meta["delegation_tag"] = delegation_tag
-        self.partition_tags[delegation_tag] = self.delegation_spec_1
+        partition_tags[delegation_tag] = self.delegation_spec_1
 
         # Tag all nodes in the first partiton to backend 2
         node_to_backend_2 = ... # some logic to select the nodes from the graph
         delegation_tag = f"backend2_tag{partitioner_2.id}"
         node.meta["delegation_tag"] = delegation_tag
-        self.partition_tags[delegation_tag] = self.delegation_spec_2
-        return edge_graph_module
+        partition_tags[delegation_tag] = self.delegation_spec_2
+        return PartitionResult(
+            tagged_graph=edge_graph_module, partition_tags=partition_tags
+        )
 ```
 
 6. Is there an easy way to write partitioner?

exir/backend/backend_api.py

@@ -16,7 +16,11 @@
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 
-from executorch.exir.backend.partitioner import Partitioner, TPartitioner
+from executorch.exir.backend.partitioner import (
+    Partitioner,
+    PartitionResult,
+    TPartitioner,
+)
 from executorch.exir.backend.utils import is_identical_graph
 
 from executorch.exir.delegate import executorch_call_delegate, get_lowered_module_name
@@ -142,10 +146,10 @@ def validation_disabled() -> Generator[None, None, None]:
 
 def _partition_and_lower(
     tagged_graph_module: torch.fx.GraphModule,
-    partitioner_instance: Partitioner,
+    partition_result: PartitionResult,
     owning_program: ExportedProgram,
 ) -> torch.fx.GraphModule:
-    for tag, delegation_spec in partitioner_instance.partition_tags.items():
+    for tag, delegation_spec in partition_result.partition_tags.items():
         # Create partition with nodes containing this tag. There should only be
         # one contained submodule per tag
         node_list = []
@@ -231,7 +235,7 @@ def _partition_and_lower(
     # Recursively partition and lower for submodules
     for name, submod, _node in get_control_flow_submodules(tagged_graph_module):
         partitioned_submodule = _partition_and_lower(
-            submod, partitioner_instance, owning_program
+            submod, partition_result, owning_program
         )
         tagged_graph_module.add_module(name, partitioned_submodule)
 
@@ -281,7 +285,8 @@ def to_backend(
     copied_graph_module = copy.deepcopy(edge_graph_module)
     # Call the partitioner on the given graph module
     partitioner_instance: Partitioner = partitioner()
-    tagged_graph_module = partitioner_instance(copied_graph_module)
+    partitioner_result = partitioner_instance(copied_graph_module)
+    tagged_graph_module = partitioner_result.tagged_graph
 
     # Check that the partitioner did not modify the original graph
     if _ENABLE_VALIDATION:
@@ -293,12 +298,11 @@ def to_backend(
         logging.warning("Disabled validating the partitioner.")
 
     assert (
-        hasattr(partitioner_instance, "partition_tags")
-        and partitioner_instance.partition_tags is not None
+        partitioner_result.partition_tags is not None
     ), f"Partitioner {partitioner} needs a `partition_tags` field containing a mapping of tags to delegate spec"
 
     tagged_graph_module = _partition_and_lower(
-        tagged_graph_module, partitioner_instance, edge_program
+        tagged_graph_module, partitioner_result, edge_program
     )
 
     # TODO(angelayi): Update this signature in a less manual way (maybe through

exir/backend/partitioner.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import Dict, List, NamedTuple, TypeVar
 
 import torch.fx as fx
@@ -18,14 +19,27 @@ class DelegationSpec(NamedTuple):
     compile_specs: List[CompileSpec]
 
 
+@dataclass
+class PartitionResult:
+    """
+    tagged_graph: the graph with nodes that intend to be delegated containing a "DelegationSpec" metadata
+    partition_tags: A dictionary that will be used to keep track of the tags and it's corresponding DelegationSpec. The tag is defined by users and used
+    in the node.meta.
+    """
+
+    tagged_graph: fx.GraphModule
+    partition_tags: Dict[str, DelegationSpec]
+
+
 class Partitioner(ABC):
     """
-    Defines a callable interface for partitioning an exported Module (i.e. a program) for
+    Defines a callable interface for partitioning an exported module (i.e. a program) for
     backend delegation.
-    A partitioner implementation would receive an exported Module, determine what portions of
+    A partitioner implementation would receive an exported module, determine what portions of
     the it can be delegated to certain backend (though a partitioner can target multiple
-    backends as well), and return the same input Module with specific nodes in
-    the input graph tagged for delegation.
+    backends as well), and return the PartitionResult including:
+    - the same input module with specific nodes in the input graph tagged for delegation
+    - the "partition_tags" to indicate how the tag is mapped to Delegation Spec.
 
     The nodes that intend to be delegated must be tagged (by setting
     node.meta["delegation_tag"]) and this tag must be provided in the
@@ -40,14 +54,12 @@ class Partitioner(ABC):
         edge_graph_module: A module in Edge dialect to be partitioned for backend delegation.
     """
 
-    partition_tags: Dict[str, DelegationSpec]
-
-    def __call__(self, edge_graph_module: fx.GraphModule) -> fx.GraphModule:
+    def __call__(self, edge_graph_module: fx.GraphModule) -> PartitionResult:
         return self.partition(edge_graph_module)
 
     @enforcedmethod
     @abstractmethod
-    def partition(self, edge_graph_module: fx.GraphModule) -> fx.GraphModule:
+    def partition(self, edge_graph_module: fx.GraphModule) -> PartitionResult:
         """
         Returns the input exported program with newly created sub-Modules encapsulating
         specific portions of the input "tagged" for delegation.
@@ -66,8 +78,7 @@ def partition(self, edge_graph_module: fx.GraphModule) -> fx.GraphModule:
             edge_graph_module: A module in Edge dialect to be partitioned for backend delegation.
 
         Returns:
-            GraphModule: Returns the input exported program with nodes that
-            intend to be delegated containing a "delegate_spec" metadata
+            PartitionResult: includes the tagged graph and the delegation spec to indicate what backend_id and compile_spec is used for each node and the tag created by the backend developers.
         """
         pass
 

exir/backend/test/demos/rpc/executor_backend_partitioner.py

@@ -11,7 +11,11 @@
 from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import (
     generate_pattern_op_partitions,
 )
-from executorch.exir.backend.partitioner import DelegationSpec, Partitioner
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
 from executorch.exir.backend.test.backend_with_compiler_demo import (
     BackendWithCompilerDemo,
 )
@@ -45,23 +49,23 @@ def __init__(self) -> None:
         self.op_support = any_chain(AnyOperatorSupport(), AnyDelegateSupport())
         self.delegation_spec = DelegationSpec("ExecutorBackend", [])
 
-        self.partition_tags = {}
-
-    def partition(
-        self, edge_graph_module: torch.fx.GraphModule
-    ) -> torch.fx.GraphModule:
+    def partition(self, edge_graph_module: torch.fx.GraphModule) -> PartitionResult:
+        partition_tags = {}
         partition_list = generate_pattern_op_partitions(
             edge_graph_module, op_support=self.op_support
         )
         for partition in partition_list:
             for node in partition.nodes:
                 delegation_tag = f"tag{partition.id}"
                 node.meta["delegation_tag"] = delegation_tag
-                self.partition_tags[delegation_tag] = self.delegation_spec
+                partition_tags[delegation_tag] = self.delegation_spec
 
                 # Tag the delegate submodules
                 # pyre-ignore Undefined attribute [16]: Item `None` of `typing.Union[None, typing.Dict[str, typing.Any], typing.List[typing.Any], bool, complex, float, int, range, slice, str, torch._C.device, torch._C.dtype, torch._C.layout, torch._C.memory_format, torch._tensor.Tensor, torch.fx.node.Node, typing.Tuple[typing.Any, ...]]` has no attribute `op`.Pyre
                 if node.args[0].op == "get_attr":
                     # pyre-ignore Undefined attribute [16]: Item `None` of `typing.Union[None, typing.Dict[str, typing.Any], typing.List[typing.Any], bool, complex, float, int, range, slice, str, torch._C.device, torch._C.dtype, torch._C.layout, torch._C.memory_format, torch._tensor.Tensor, torch.fx.node.Node, typing.Tuple[typing.Any, ...]]` has no attribute `op`.Pyre
                     node.args[0].meta["delegation_tag"] = delegation_tag
-        return edge_graph_module
+
+        return PartitionResult(
+            tagged_graph=edge_graph_module, partition_tags=partition_tags
+        )