New util function to print delegation summary

Summary:
Because what happened to the graph in to_backend() is pretty opaque to the users, the existing [print_delegated_graph()](https://fburl.com/code/4xa5oewv) is good but might be too much for new users, so implemented this new util function to give a more lightweight summary for the delegated graph. 
Original design doc:
https://docs.google.com/document/d/19ZSDddm23MnGvFUrkV9clwHwLvXzAo2IZtbkr9VChjE/edit?usp=sharing

**Note:** The added `import pandas` shouldn't be an issue because we already have pandas in the install requirement is executorch oss: https://fburl.com/code/p55654r2

Differential Revision: D55239751

Author: Olivia-liu

exir/backend/test/test_utils.py

@@ -14,6 +14,7 @@
 from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo
 from executorch.exir.backend.utils import (
     get_delegates,
+    get_delegation_info,
     get_non_lowered_nodes,
     is_identical_graph,
     print_delegated_graph,
@@ -439,3 +440,47 @@ def forward(self, a, x, b):
             graph_str,
             "Expect to see the aten.mm in the delegated graph",
         )
+
+    def test_get_delegation_info(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, a, x, b):
+                y = torch.mm(a, x)
+                z = y + b
+                a = z - a
+                y = torch.mm(a, x)
+                z = y + b
+                return z
+
+        m = Model()
+        inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
+
+        edge = to_edge(torch.export.export(m, inputs)).to_backend(
+            AddMulPartitionerDemo()
+        )
+
+        delegation_info = get_delegation_info(edge.exported_program().graph_module)
+        self.assertEqual(delegation_info.num_delegated_nodes, 4)
+        self.assertIn(
+            "Total delegated subgraphs",
+            delegation_info.get_summary(),
+        )
+        df = delegation_info.get_operator_delegation_dataframe()
+        self.assertFalse(df.empty)
+        integer_columns = [
+            "occurrences in delegated graphs",
+            "occurrences in non-delegated graphs",
+        ]
+        has_positive = False
+        for column in integer_columns:
+            # Check if there is at least one positive number in the column
+            if df[column].gt(0).any():
+                has_positive = True
+                break
+
+        self.assertTrue(
+            has_positive,
+            f"The column '{column}' does not contain any positive numbers",
+        )

exir/backend/utils.py

@@ -6,10 +6,13 @@
 
 import logging
 import operator
+import re
 from collections import defaultdict
+from dataclasses import dataclass
 from functools import lru_cache
 from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
 
+import pandas as pd
 import torch
 from executorch.exir.backend.backend_details import ExportedProgram
 from executorch.exir.common import setting_python_recursive_limit
@@ -211,6 +214,130 @@ def get_delegates(graph: torch.fx.Graph) -> List[torch.fx.Node]:
     ]
 
 
+@dataclass
+class DelegationBreakdown:
+    delegated: int
+    non_delegated: int
+
+
+@dataclass
+class DelegationInfo:
+    num_delegated_subgraphs: int
+    num_delegated_nodes: int
+    num_non_delegated_nodes: int
+    delegation_by_operator: Dict[str, DelegationBreakdown]
+
+    def get_summary(self) -> str:
+        # Assemble and return the summary string
+        summary_str = f"Total delegated subgraphs: {self.num_delegated_subgraphs}\n"
+        summary_str += f"Number of delegated nodes: {self.num_delegated_nodes}\n"
+        summary_str += (
+            f"Number of non-delegated nodes: {self.num_non_delegated_nodes}\n"
+        )
+        return summary_str
+
+    def get_operator_delegation_dataframe(self) -> pd.DataFrame:
+        # Convert the dict to a dataframe
+        df = pd.DataFrame(
+            list(self.delegation_by_operator.items()),
+            columns=["op type", "occurrences"],
+        )
+
+        # Function to extract delegated and non-delegated fields
+        def _extract_breakdown(row):
+            return pd.Series(
+                [row["occurrences"].delegated, row["occurrences"].non_delegated]
+            )
+
+        df[
+            ["occurrences in delegated graphs", "occurrences in non-delegated graphs"]
+        ] = df.apply(lambda row: _extract_breakdown(row), axis=1)
+        df.drop(columns=["occurrences"], inplace=True)
+
+        # Add a Total row at the bottom
+        total_delegated_nodes = df["occurrences in delegated graphs"].sum()
+        total_non_delegated_nodes = df["occurrences in non-delegated graphs"].sum()
+        df = df.sort_values(by="op type", ignore_index=True)
+        df.loc["Total"] = ["Total", total_delegated_nodes, total_non_delegated_nodes]
+
+        return df
+
+
+def _get_op_type(node_name: str) -> str:
+    # node_name is in format <op_type> or <op_type>_x in which x is an integer suffix.
+    pattern = r"_(\d+)$"
+    match = re.search(pattern, node_name)
+    if match:
+        return node_name[: match.start()]
+    else:
+        return node_name
+
+
+def get_delegation_info(
+    graph_module: torch.fx.GraphModule,
+) -> DelegationInfo:
+    """
+    Returns a string (for high level summary) and a DataFrame (for per op type summary) of the
+    delegation info on the given graph module.
+    """
+
+    op_occurrences_dict = {}
+
+    def _insert_op_occurrences_dict(node_name: str, delegated: bool) -> None:
+        op_type = _get_op_type(node_name)
+        if op_type not in op_occurrences_dict:
+            if delegated:
+                op_occurrences_dict[op_type] = DelegationBreakdown(
+                    delegated=1, non_delegated=0
+                )
+            else:
+                op_occurrences_dict[op_type] = DelegationBreakdown(
+                    delegated=0, non_delegated=1
+                )
+        else:
+            if delegated:
+                op_occurrences_dict[op_type].delegated += 1
+            else:
+                op_occurrences_dict[op_type].non_delegated += 1
+
+    delegated_subgraph_counter = 0
+
+    lowered_module_dict = {
+        node.name: getattr(graph_module, node.name)
+        for node in graph_module.graph.nodes
+        if node.op == "get_attr" and node.name.startswith("lowered_module_")
+    }
+
+    for node in graph_module.graph.nodes:
+        if (
+            node.op == "call_function"
+            and _get_op_type(node.name) != "executorch_call_delegate"
+        ):
+            _insert_op_occurrences_dict(node_name=node.name, delegated=False)
+        if node.op == "get_attr" and node.name.startswith("lowered_module_"):
+            lowered_module = lowered_module_dict[node.name]
+            delegated_subgraph_counter += 1
+            for node_in_lowered_module in lowered_module.original_module.graph.nodes:
+                if node_in_lowered_module.op == "call_function":
+                    _insert_op_occurrences_dict(
+                        node_name=node_in_lowered_module.name, delegated=True
+                    )
+
+    # Calculate the total number of delegated and non-delegated nodes
+    num_delegated_nodes = 0
+    num_non_delegated_nodes = 0
+    for _, value in op_occurrences_dict.items():
+        num_delegated_nodes += value.delegated
+        num_non_delegated_nodes += value.non_delegated
+
+    return DelegationInfo(
+        num_delegated_nodes=num_delegated_nodes,
+        num_non_delegated_nodes=num_non_delegated_nodes,
+        num_delegated_subgraphs=delegated_subgraph_counter,
+        delegation_by_operator=op_occurrences_dict,
+    )
+
+
 def print_delegated_graph(graph_module: torch.fx.GraphModule) -> str:
     """
     Print the graph of including lowered_module (both backend id and original graph) together with the graph module. Example output: