Introduce DelegateMappingBuilder (#184)

Jack-Khuu · facebook-github-bot · commit 9ca4ed025b09 · 2023-09-01T02:29:35.000-07:00
Summary: Pull Request resolved: #184 For creating the debug_handle_map generated in the design doc, this diff introduces a helper class for manage the state of the handle map `get_delegate_mapping`: Is used to return the debug handle map out of preprocess `upsert_delegate_mapping_entry`: Handles the creation and manipulation of debug handle entries --- Reviewed By: tarun292 Differential Revision: D48569900 fbshipit-source-id: 2a6e462d9306380465a46ae20b2aeb0cb1c4dc67
diff --git a/exir/backend/test/TARGETS b/exir/backend/test/TARGETS
@@ -211,6 +211,18 @@ python_unittest(
     ],
 )
 
+python_unittest(
+    name = "test_delegate_map_builder",
+    srcs = [
+        "test_delegate_map_builder.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/exir/backend:utils",
+    ],
+)
+
 python_unittest(
     name = "test_utils",
     srcs = [
diff --git a/exir/backend/test/test_delegate_map_builder.py b/exir/backend/test/test_delegate_map_builder.py
@@ -0,0 +1,167 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from typing import Iterator, Union
+
+import torch
+from executorch import exir
+
+from executorch.exir.backend.utils import DelegateMappingBuilder
+
+
+class TestDelegateMapBuilder(unittest.TestCase):
+    def setUp(self) -> None:
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                y = torch.sin(x)
+                return torch.cos(y)
+
+        model = Model()
+        model_inputs = (torch.ones(1, 1),)
+        program = (
+            exir.capture(model, model_inputs, exir.CaptureConfig(pt2_mode=True))
+            .to_edge()
+            .to_executorch()
+        )
+
+        # Create nodes for testing mapping
+        # nodes: [arg0_1, alloc, aten_sin_default, alloc_1, aten_cos_default, output]
+        # debug handles: [0, None, 1, None, 2, 3]
+        self.nodes = list(program.graph_module.graph.nodes)
+
+    def test_basic_generated_identifier(self):
+        delegate_builder = DelegateMappingBuilder(generated_identifiers=True)
+
+        expected_mapping = {0: (0, 1, 2, 3)}
+        self.assertEqual(delegate_builder.upsert_delegate_mapping_entry(self.nodes), 0)
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+
+        expected_mapping = {0: (0, 1, 2, 3), 1: (0,)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes[0]), 1
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+
+    def test_basic_manual_int_identifier(self):
+        self._test_basic_manual_identifier(iter([22, 55]))
+
+    def test_basic_manual_string_identifier(self):
+        self._test_basic_manual_identifier(iter(["22", "55"]))
+
+    def test_appending_nodes_generated_identifier(self):
+        delegate_builder = DelegateMappingBuilder(generated_identifiers=True)
+
+        expected_mapping = {0: (0,)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes[0]), 0
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+        self._test_appending_nodes(delegate_builder, 0)
+
+    def test_appending_nodes_manual_int_identifier(self):
+        delegate_builder = DelegateMappingBuilder()
+
+        expected_mapping = {22: (0,)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes[0], 22), 22
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+        self._test_appending_nodes(delegate_builder, 22)
+
+    def test_appending_nodes_manual_string_identifier(self):
+        delegate_builder = DelegateMappingBuilder()
+
+        expected_mapping = {"22": (0,)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes[0], "22"), "22"
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+        self._test_appending_nodes(delegate_builder, "22")
+
+    def test_adding_manual_identifier_when_generated(self):
+        delegate_builder = DelegateMappingBuilder(generated_identifiers=True)
+
+        self.assertRaises(
+            Exception,
+            lambda: delegate_builder.upsert_delegate_mapping_entry(self.nodes, "22"),
+        )
+
+    def test_omitting_identifier_when_not_generated(self):
+        delegate_builder = DelegateMappingBuilder()
+
+        self.assertRaises(
+            Exception,
+            lambda: delegate_builder.upsert_delegate_mapping_entry(self.nodes),
+        )
+
+    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+    def _test_basic_manual_identifier(self, identifiers: Iterator[Union[int, str]]):
+        """
+        Using the iteration of identifiers:
+        1) Create a Delegate Map Builder
+        2) Add an entry with a list of Nodes using the first identifier
+        3) Add an entry with a single node using the second identifier
+
+        Verify behavior results
+        """
+        delegate_builder = DelegateMappingBuilder()
+
+        # Entry with a list of nodes
+        iden_1 = next(identifiers)
+        expected_mapping = {iden_1: (0, 1, 2, 3)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes, iden_1), iden_1
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+
+        # Entry with a single node
+        iden_2 = next(identifiers)
+        expected_mapping = {iden_1: (0, 1, 2, 3), iden_2: (0,)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes[0], iden_2),
+            iden_2,
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+
+    def _test_appending_nodes(
+        self, delegate_builder: DelegateMappingBuilder, identifier: Union[int, str]
+    ):
+        """
+        For the provided delegate_builder and identifier:
+        1) Append a duplicate Node (previously in builder)
+        2) Append a list of Nodes
+        3) Append a single Node
+
+        Verify behavior results
+        """
+        # Add a duplicate node
+        expected_mapping = {identifier: (0,)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes[0], identifier),
+            identifier,
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+
+        # Add a list of nodes
+        expected_mapping = {identifier: (0, 1, 2)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes[1:5], identifier),
+            identifier,
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
+
+        # Add a single node
+        expected_mapping = {identifier: (0, 1, 2, 3)}
+        self.assertEqual(
+            delegate_builder.upsert_delegate_mapping_entry(self.nodes[5], identifier),
+            identifier,
+        )
+        self.assertEqual(delegate_builder.get_delegate_mapping(), expected_mapping)
diff --git a/exir/backend/utils.py b/exir/backend/utils.py
@@ -4,18 +4,23 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
+from collections import defaultdict
 from functools import lru_cache
-from typing import Iterable, List, Tuple
+from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import torch
 from executorch.exir.dialects._ops import ops as exir_ops
 
 from executorch.exir.lowered_backend_module import create_submodule_from_nodes
+from torch.fx.node import Node
 from torch.fx.passes.utils.source_matcher_utils import SourcePartition
 
 T_QuantPerTensor = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
 T_DQuantPerTensor = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
 
+log: logging.Logger = logging.getLogger(__name__)
+
 
 @lru_cache(maxsize=128)
 def is_same_node(
@@ -158,3 +163,100 @@ def replace_quantized_partition_with_op(
     graph_module.recompile()
 
     return (replaced_op, dequant_nodes, quant_nodes)
+
+
+# TODO - style: use templated types
+class DelegateMappingBuilder:
+    """
+    Profiling helper class for building Delegate Maps.
+    Delegate Maps are a mapping from log entry identifiers to node
+    debug handles. Specifically this is used to log within backend delegates
+
+    Args:
+        generated_identifiers (bool, optional): Whether delegate map keys are
+            generated automatically. Defaults to False.
+    """
+
+    def __init__(self, generated_identifiers: bool = False):
+        self._generated_identifiers = generated_identifiers
+
+        # Note that the internal struct has a Set value, while the getter
+        # function returns the values as a tuple
+        self._debug_handle_map: Union[
+            Dict[int, Set[int]], Dict[str, Set[int]]
+        ] = defaultdict(set)
+        self._next_index: int = 0
+
+    def get_delegate_mapping(
+        self,
+    ) -> Union[Dict[int, Tuple[int]], Dict[str, Tuple[int]]]:
+        """
+        Returns:
+           Union[Dict[int, Tuple[int]], Dict[str, Tuple[int]]]:
+                A map of delegate debug identifier to a list of debug handles
+                The keys (identifier) are either integers or strings
+                The values are a sorted tuple of integer debug handles
+        """
+        # pyre-ignore Warning between Union[Dict[K, V], Dict[K2, V]] vs Dict[Union[K, K2], V]
+        return {k: tuple(sorted(v)) for k, v in self._debug_handle_map.items()}
+
+    def upsert_delegate_mapping_entry(
+        self,
+        nodes: Union[Node, List[Node]],
+        identifier: Optional[Union[int, str]] = None,
+    ) -> Union[int, str]:
+        """
+        Add or append to an existing delegate mapping entry
+
+        If self._generated_identifiers = False:
+            - An identifier must be provided, else this is no-op
+
+        If self._generated_identifiers = True:
+            - NEW identifiers cannot be manually provided. Existing identifier
+                can be provided
+            - New identifiers are generated incrementally, 0 indexed
+
+        If a provided identifier already exists, node entries are appended
+
+        Args:
+            nodes (Union[Node, List[Node]]): A (list of) Node(s)
+            identifier (Optional[Union[int, str]]):
+                Debug identifier corresponding to the Node(s)
+
+        Returns:
+            Optional[Union[int, str]]:
+                Delegate debug identifier corresponding to the node group
+                    None is returned, if no identifier is associated
+        """
+
+        # Check for manual addition of identifier (with generated identifiers enabled)
+        if (
+            self._generated_identifiers
+            and identifier is not None
+            and identifier not in self._debug_handle_map
+        ):
+            raise Exception(
+                f"Builders using generated identifiers can't manually add identifiers: {identifier}. Failed to add or update entry"
+            )
+
+        # Resolve Identifier
+        if identifier is None:
+            if self._generated_identifiers:
+                identifier = self._next_index
+                self._next_index += 1
+            else:
+                raise Exception(
+                    "No identifier provided. Failed to add or update entry."
+                )
+
+        # Get all debug handles found in the nodes
+        # Note that missing debug handles are not surfaced
+        new_debug_handles = {
+            handle
+            for node in (nodes if isinstance(nodes, List) else [nodes])
+            if (handle := node.meta.get("debug_handle")) is not None
+        }
+
+        # pyre-ignore Warning from Union[int, st] keys
+        self._debug_handle_map[identifier].update(new_debug_handles)
+        return identifier
