executorch 1.0.0__cp310-cp310-win_amd64.whl

executorch/exir/backend/utils.py

@@ -0,0 +1,661 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+import logging
+import operator
+from collections import defaultdict, OrderedDict
+from functools import lru_cache
+from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
+
+import torch
+from executorch.exir.backend.backend_details import ExportedProgram
+from executorch.exir.backend.canonical_partitioners.duplicate_constant_node_pass import (
+    duplicate_constant_node,
+)
+from executorch.exir.common import setting_python_recursive_limit
+from executorch.exir.delegate import executorch_call_delegate
+from executorch.exir.dialects._ops import ops as exir_ops
+
+from executorch.exir.lowered_backend_module import create_submodule_from_nodes
+from tabulate import tabulate
+from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
+from torch.fx.experimental.symbolic_shapes import has_free_symbols
+from torch.fx.node import Node
+from torch.fx.passes.operator_support import OperatorSupportBase
+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
+
+
+# NB: Set this to None to handle validation from MobileBert
+@lru_cache(maxsize=None)
+def is_same_node(
+    node_left: Iterable[torch.fx.Node],
+    node_right: Iterable[torch.fx.Node],
+) -> bool:
+    # two nodes are the same if they have the same target and op
+    # same for their args
+    if isinstance(node_left, torch.fx.Node) and isinstance(node_right, torch.fx.Node):
+        if not (
+            (node_left.target == node_right.target)
+            and (node_left.op == node_right.op)
+            and (len(node_left.all_input_nodes) == len(node_right.all_input_nodes))
+            and all(
+                is_same_node(arg_left, arg_right)
+                for arg_left, arg_right in zip(
+                    node_left.all_input_nodes, node_right.all_input_nodes
+                )
+            )
+        ):
+            return False
+    else:
+        if len(list(node_left)) != len(list(node_right)):
+            return False
+        for n_left, n_right in zip(node_left, node_right):
+            if not is_same_node(n_left, n_right):
+                return False
+    return True
+
+
+def is_identical_graph(
+    graph_left: torch.fx.GraphModule, graph_right: torch.fx.GraphModule
+) -> bool:
+    # two graph are the same if they have the same nodes and op. The order of nodes also
+    # matters in this function is more strict. Two graph are not considered as the same
+    # if the topological order of the nodes is the same in this function but the order of nodes
+    # is not the same.
+    if len(list(graph_left.graph.nodes)) != len(list(graph_right.graph.nodes)):
+        return False
+    with setting_python_recursive_limit(30000):
+        for node_left, node_right in zip(
+            graph_left.graph.nodes, graph_right.graph.nodes
+        ):
+            if not (is_same_node(node_left, node_right)):
+                return False
+    return True
+
+
+def remove_first_quant_and_last_dequant(
+    graph_module: torch.fx.GraphModule,
+) -> None:
+    for node in graph_module.graph.nodes:
+        if node.target == T_QuantPerTensor:
+            if node.args[0].op == "placeholder":
+                node_users = list(node.users.keys())
+                for dequant_node in node_users:
+                    # point the dequant arg to the placeholder
+                    dequant_node.args = (node.args[0],) + dequant_node.args[1:]
+        elif node.target == T_DQuantPerTensor:
+            node_users = list(node.users.keys())
+            if node_users[0].op == "output":
+                # point the output arg to the quant node
+                output_node = node_users[0]
+                output_node.args = ([node.args[0]],)
+    # Remove the quant/dequant nodes as they don't have users
+    graph_module.graph.eliminate_dead_code()
+    graph_module.recompile()
+
+
+# TODO - use edge ops
+def replace_quantized_partition_with_op(
+    graph_module: torch.fx.GraphModule,
+    partition: SourcePartition,
+    replacement_op: torch._ops.OpOverloadPacket,
+) -> Tuple[torch.fx.Node, List[torch.fx.Node], List[torch.fx.Node]]:
+    """
+    Replaces partition with the op specified by replacement_op. It's also expected that
+    the nodes contained in partition are sourced from a quantized module as this function
+    searches for the quantization pattern to consume along with the nodes in the partition,
+    to be then replaced by replacement_op.
+
+    Args:
+        graph_module: The graph module from which this partition was sourced.
+        partition: Partition to be replaced.
+        replacement_op: The op to replace paritition with.
+    Returns:
+        Tuple: First element in the tuple is the new replaced module. The second and third
+        node lists in the returned tuple consist of the dq and q nodes that were consumed
+        along with this partition to be replaced by the replacement_op.
+    """
+
+    dequant_nodes = []
+    quant_nodes = []
+    input_nodes = []
+    output_nodes = []
+
+    partition_nodes = [node for node in partition.nodes if node not in partition.params]
+
+    # We recreate our input nodes and output nodes list instead of using partition.input_nodes
+    # and partition.output_nodes as the ordering of the nodes in those lists is not deterministic,
+    # whereas for the quant fusion pass we expect deterministic ordering.
+    for node in partition.nodes:
+        for arg in node.args:
+            if isinstance(arg, torch.fx.Node) and (arg not in partition.nodes):
+                input_nodes.append(arg)
+
+        for user in node.users.keys():
+            if user not in partition.nodes:
+                output_nodes.append(node)
+
+    # Try to find all the dq nodes that are feeding into this module.
+    for node in input_nodes:
+        if node.target == T_DQuantPerTensor:
+            dequant_nodes += [node]
+
+    # Try to find all the q nodes that this module is feeding out into.
+    for node in output_nodes:
+        for user in node.users.keys():
+            if user.target == T_QuantPerTensor:
+                quant_nodes += [user]
+
+    assert len(dequant_nodes) >= 1, "Dequant nodes missing in node list to be replaced."
+    assert len(quant_nodes) >= 1, "Quant nodes missing in node list to be replaced."
+
+    # After this, node list will essentially contain all the nodes in the
+    # dq->op->q pattern that we will want to replace with a custom backend op.
+    node_list = dequant_nodes + partition_nodes + quant_nodes
+
+    submodule, call_module_node = create_submodule_from_nodes(
+        graph_module, node_list, "to_be_replaced", skip_legalize_graph=True
+    )
+
+    # Update the replaced op so that we have all the latest args and kwargs.
+    with graph_module.graph.inserting_before(call_module_node):
+        replaced_op = graph_module.graph.call_function(
+            replacement_op,
+            call_module_node.args,
+            kwargs=call_module_node.kwargs,
+        )
+        call_module_node.replace_all_uses_with(replaced_op)
+        graph_module.graph.erase_node(call_module_node)
+        replaced_op.meta = call_module_node.meta
+    graph_module.recompile()
+
+    return (replaced_op, dequant_nodes, quant_nodes)
+
+
+def _assign_new_tag(
+    tagged_exported_program: ExportedProgram,
+    copied_nodes: Set[str],
+):
+    """
+    Assign new tag to the copied nodes.
+
+    Before the pass
+    constant_0 (tag_10) ------------------> op_b (tag_10)
+    constant_0_copy (tag_10) -------------> op_a (tag_11)
+
+    After the pass
+    constant_0 (tag_10) ------------------> op_b (tag_10)
+    constant_0_copy (tag_11) -------------> op_a (tag_11)
+
+    """
+    for node in tagged_exported_program.graph.nodes:
+        if node.op == "placeholder":
+            if node.name in copied_nodes:
+                users_tag = set()
+                for user in node.users:
+                    users_tag.add(user.meta.get("delegation_tag", None))
+                # Assign the tag to the copy constant node the same as their users.
+                if len(users_tag) == 1:
+                    node.meta["delegation_tag"] = users_tag.pop()
+
+
+def _maybe_duplicate_constant_nodes(
+    tagged_exported_program: ExportedProgram,
+    tag: str,
+) -> None:
+    """
+    If the constants node is shared by different tagged nodes, like
+    constant_0 ----> op_b (tag_10)
+    |-------------> op_a (tag_11)
+
+    we make default as constant_0 is duplicated to constant_0_1, constant_0_2, unless the node is tagged with "no_copy"
+    constant_0 ------------------> op_b (tag_10)
+    constant_0_copy -------------> op_a (tag_11)
+
+    backend can estimate how much they want to duplicate the constant node, either error out or default to duplicate
+    """
+    candidate_nodes = set()
+    for node in tagged_exported_program.graph.nodes:
+        if node.meta.get("delegation_tag", "") == tag:
+            if node.op == "placeholder":
+                for user in node.users:
+                    users_tag = user.meta.get("delegation_tag", None)
+                    if users_tag != tag:
+                        # If the node is tagged with "no_copy", we stop duplicating it and throw an error
+                        if node.meta.get("no_copy", False):
+                            raise RuntimeError(
+                                f"constant data node ({node}) is tagged with ({tag}) but has user ({user}) which has tag ({users_tag})"
+                            )
+                        else:
+                            candidate_nodes.add(node.name)
+    copied_nodes = set()
+    for candidate_node in candidate_nodes:
+        # Both tagged exported program and the owning program need to go through the same duplication pass
+        copied_nodes = copied_nodes.union(
+            duplicate_constant_node(tagged_exported_program, candidate_node)
+        )
+    candidate_node_with_copies = candidate_nodes.union(copied_nodes)
+    _assign_new_tag(tagged_exported_program, candidate_node_with_copies)
+
+
+def _get_item_from_executorch_call_delegate(node: torch.fx.Node) -> bool:
+    """
+    Check if the node is the getitem followed by executorch_call_delegate node. These getitems node
+    are just for getting the result from delegate because the input/output to delegates are flattened
+    """
+    return (
+        node.target == operator.getitem
+        and len(node.args) == 2
+        and node.args[0].target == executorch_call_delegate  # pyre-ignore
+        and isinstance(node.args[1], int)
+    )
+
+
+def get_non_lowered_nodes(graph: torch.fx.Graph) -> List[torch.fx.Node]:
+    """
+    Returns a list of non lowered nodes in the graph module.
+    """
+    return [
+        node
+        for node in graph.nodes
+        if node.op == "call_function"
+        and node.target != executorch_call_delegate
+        and (not _get_item_from_executorch_call_delegate(node))
+    ]
+
+
+def get_delegates(graph: torch.fx.Graph) -> List[torch.fx.Node]:
+    """
+    Returns the list of delegates from the graph.
+    """
+    return [
+        node
+        for node in graph.nodes
+        if node.op == "get_attr" and node.name.startswith("lowered_module_")
+    ]
+
+
+def print_delegated_graph(graph_module: torch.fx.GraphModule) -> None:
+    """
+    Print the formatted graph string.
+    """
+    print(format_delegated_graph(graph_module))
+
+
+def format_delegated_graph(graph_module: torch.fx.GraphModule) -> str:
+    """
+    Return the formatted graph string of including lowered_module (both backend id and original graph) together with the graph module. Example output:
+    graph():
+        %arg0_1 : [num_users=2] = placeholder[target=arg0_1]
+        %arg1_1 : [num_users=2] = placeholder[target=arg1_1]
+        %arg2_1 : [num_users=2] = placeholder[target=arg2_1]
+        %lowered_module_0 : [num_users=1] = get_attr[target=lowered_module_0]
+            backend_id: BackendWithCompilerDemo
+            lowered graph():
+                %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
+                %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
+                %arg2_1 : [num_users=1] = placeholder[target=arg2_1]
+                %aten_mm_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.mm.default](args = (%arg0_1, %arg1_1), kwargs = {})
+                %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%aten_mm_default, %arg2_1), kwargs = {})
+                return [aten_add_tensor]
+        %executorch_call_delegate : [num_users=1] = call_function[target=torch.ops.higher_order.executorch_call_delegate](args = (%lowered_module_0, %arg0_1, %arg1_1, %arg2_1), kwargs = {})
+        %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%executorch_call_delegate, 0), kwargs = {})
+        %aten_sub_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.sub.Tensor](args = (%getitem, %arg0_1), kwargs = {})
+        %lowered_module_1 : [num_users=1] = get_attr[target=lowered_module_1]
+            backend_id: BackendWithCompilerDemo
+            lowered graph():
+                %aten_sub_tensor : [num_users=1] = placeholder[target=aten_sub_tensor]
+                %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
+                %arg2_1 : [num_users=1] = placeholder[target=arg2_1]
+                %aten_mm_default_1 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.mm.default](args = (%aten_sub_tensor, %arg1_1), kwargs = {})
+                %aten_add_tensor_1 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%aten_mm_default_1, %arg2_1), kwargs = {})
+                return [aten_add_tensor_1]
+        %executorch_call_delegate_1 : [num_users=1] = call_function[target=torch.ops.higher_order.executorch_call_delegate](args = (%lowered_module_1, %aten_sub_tensor, %arg1_1, %arg2_1), kwargs = {})
+        %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%executorch_call_delegate_1, 0), kwargs = {})
+        return [getitem_1]
+    """
+    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_")
+    }
+    indent = "  "
+    graph_format_str = "graph():\n"
+    for node in graph_module.graph.nodes:
+        graph_format_str += f"{indent}{node.format_node()}\n"
+        if node.op == "get_attr" and node.name.startswith("lowered_module_"):
+            lowered_module = lowered_module_dict[node.name]
+            graph_format_str += f"{indent * 2}backend_id: {lowered_module.backend_id}\n"
+            graph_format_str += f"{indent * 2}lowered graph():\n"
+            for node_in_lowered_module in lowered_module.original_module.graph.nodes:
+                graph_format_str += (
+                    f"{indent * 3}{node_in_lowered_module.format_node()}\n"
+                )
+    return graph_format_str
+
+
+def tag_constant_data(edge_program: ExportedProgram) -> None:
+    """
+    Util function for partitioners. This function tags the const/param/buffers nodes
+    whose users all belong within the same partition. This should be called after tagging all other nodes.
+    Any const/param/buffer which is used as input to a subgraph, will be tagged with the same tag as that
+    subgraph. Throw error when const/param/buffers is used across different partitions. That is the
+    underlying data will be owned by multiple delegates.
+    """
+    mutated_buffer = set()
+    for node in edge_program.graph.nodes:
+        if node.op == "placeholder" and (
+            is_param(edge_program, node)
+            or is_buffer(edge_program, node)
+            or is_lifted_tensor_constant(edge_program, node)
+        ):
+            for node_user in node.users:
+                if node_user.name in edge_program.graph_signature.buffers_to_mutate:
+                    logging.info(
+                        "The buffer node is a mutated buffer node, which is not constant."
+                    )
+                    mutated_buffer.add(node)
+
+    for node in edge_program.graph.nodes:
+        # go through const/param/buffer nodes, if all users of const/param/buffer nodes are partitioned then partition
+        if node.op == "placeholder" and (
+            is_param(edge_program, node)
+            or is_buffer(edge_program, node)
+            or is_lifted_tensor_constant(edge_program, node)
+        ):
+            if node not in mutated_buffer:
+                user_tags = set()
+                for user in node.users:
+                    user_tag = user.meta.get("delegation_tag", None)
+                    if user_tag is not None:
+                        user_tags.add(user_tag)
+                if len(user_tags) > 1:
+                    logging.info(
+                        f"The data node is used across multiple partitions, including {user_tags}. "
+                        "If the data is too large and it's not preferred to copy, please tag the "
+                        "constant node like node.['no_copy'] = True and they won't be copied."
+                    )
+                # tag the data node with the same tag as the last user
+                if len(user_tags) > 0:
+                    node.meta["delegation_tag"] = user_tags.pop()
+
+
+def tag_mutated_buffer(edge_program: ExportedProgram) -> None:
+    """
+    Util function for partitioners. This function tags the mutated buffer nodes
+    whose users all belong within the same partition. This should be called after tagging all other nodes.
+    Any buffer which is used as input to a subgraph, will be tagged with the same tag as that
+    subgraph. Throw error when buffers is used across different partitions. That is the
+    underlying data will be owned by multiple delegates.
+    """
+    for node in edge_program.graph.nodes:
+        # Determine whether this node is a mutated buffer
+        is_mutated_buffer_node = False
+        if node.op == "placeholder" and is_buffer(edge_program, node):
+            for node_user in node.users:
+                if node_user.name in edge_program.graph_signature.buffers_to_mutate:
+                    is_mutated_buffer_node = True
+                    break
+        # This node is mutated buffer, tag it
+        if is_mutated_buffer_node:
+            user_tags = set()
+            for user in node.users:
+                user_tag = user.meta.get("delegation_tag", None)
+                if user_tag is not None:
+                    user_tags.add(user_tag)
+            if len(user_tags) > 1:
+                logging.info(
+                    f"The data node is used across multiple partitions, including {user_tags}. "
+                    "If the data is too large and it's not preferred to copy, please tag the "
+                    "constant node like node.['no_copy'] = True and they won't be copied."
+                )
+            # tag the data node with the same tag as the last user
+            if len(user_tags) > 0:
+                node.meta["delegation_tag"] = user_tags.pop()
+
+
+def is_shape_dynamic(node: torch.fx.Node) -> bool:
+    """
+    Check if the node shape is dynamic.
+    """
+
+    return has_free_symbols(node.meta["val"].shape)
+
+
+# TODO - style: use templated types
+class DelegateMappingBuilder:
+    """
+    Profiling helper class for building Delegate Mappings.
+    Delegate Mappings are mappings from delegate debug identifiers to node
+    debug handles. Specifically this is used to log within backend delegates
+
+    Args:
+        generated_identifiers (bool, optional): Whether identifier 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 insert_delegate_mapping_entry(
+        self,
+        nodes: Optional[Union[Node, List[Node]]] = None,
+        handles: Optional[Union[int, List[Optional[int]]]] = None,
+        identifier: Optional[Union[int, str]] = None,
+    ) -> Union[int, str]:
+        """
+        Add a new delegate mapping entry
+
+        If self._generated_identifiers = False:
+            - A new identifier must be provided, else an exception is thrown
+
+        If self._generated_identifiers = True:
+            - New identifiers are generated incrementally, 0 indexed
+            - Identifiers cannot be manually provided, else an exception is thrown
+
+        Args:
+            nodes (Union[Node, List[Node]]): A (list of) Node(s)
+            handles (Union[int, List[Optional[int]]]): A (list of) debug handle(s)
+            identifier (Optional[Union[int, str]]):
+                Debug identifier corresponding to the Node(s)
+
+        Note: Exactly one of nodes and handles must be provided
+        Note: If a debug handle is missing or None, it is skipped
+
+        Returns:
+            Union[int, str]:
+                Delegate debug identifier inserted
+        """
+
+        # Check for manual addition of identifier (with generated identifiers enabled)
+        if self._generated_identifiers and identifier is not None:
+            raise Exception(
+                f"Builders using generated identifiers can't manually add identifiers: {identifier}. Failed to add or update entry"
+            )
+
+        if identifier is not None and identifier in self._debug_handle_map:
+            raise Exception(
+                "This delegate debug identifier was already inserted. Duplicate delegate debug identifiers are not allowed."
+            )
+
+        # Check for exactly one of nodes and handles being populated
+        if not ((nodes is not None) ^ (handles is not None)):
+            raise Exception(
+                "Only one of nodes or handles must be provided. Either both were provided or neither were provided. 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."
+                )
+
+        # Collect debug handles
+        if nodes is not None:
+            new_debug_handles = {
+                node.meta.get("debug_handle")
+                for node in (nodes if isinstance(nodes, List) else [nodes])
+            }
+        else:
+            new_debug_handles = (
+                handles if isinstance(handles, (tuple, List)) else [handles]
+            )
+
+        # Filter for empty debug handles
+        filtered_debug_handles = {
+            handle for handle in new_debug_handles if handle is not None
+        }
+        if len(filtered_debug_handles) == 0:
+            raise Exception("No valid debug handles found. Failed to add entry.")
+
+        # pyre-ignore Warning from Union[int, st] keys
+        self._debug_handle_map[identifier] = filtered_debug_handles
+        return identifier
+
+
+class WhyNoPartition:
+    """
+    Simple helper class for partitioners to log why a node was not lowered.
+
+    Example usage:
+
+        # In your backend partitioner file(s)
+        why = WhyNoPartition(logger=your_backend_logger)
+
+        # hypothetical function that checks if a node can be lowered
+        if not can_be_lowered(node):
+            why(node, "This node was not lowered because ...")
+    """
+
+    def __init__(self, logger: logging.Logger):
+        self.logger = logger
+        self.node: Optional[torch.fx.Node] = None
+        self.reason: str = ""
+
+    def __call__(self, node: torch.fx.Node, reason: str) -> None:
+        self.node = node
+        self.reason = reason
+        self.logger.debug(self)
+
+    def __str__(self) -> str:
+        return f"WhyNoPartition: Node {self.node} was not partitioned because {self.reason}."
+
+
+class WhyNoPartitionReporter:
+    """
+    Helper class for partitioners to gather why nodes were not lowered in a single report.
+    If a node is reported multiple times, only the first report is included.
+
+    Example usage:
+
+        # In your backend partitioner file(s)
+        reporter = WhyNoPartitionReporter()
+
+        # hypothetical function that checks if a node can be lowered
+        if not can_be_lowered(node):
+            reporter.report_reject(node, "This node was not lowered because ...")
+
+        # Back in partitioner
+        logger.info(reporter.get_table_report())
+    """
+
+    def __init__(self):
+        self._rejected_nodes: OrderedDict[torch.fx.Node, str] = (
+            OrderedDict()
+        )  # {Rejected node: reason}
+
+    def report_reject(self, node: torch.fx.Node, reason: str):
+        """Report a node that was rejected from a partition, along with a reason for why."""
+        if node not in self._rejected_nodes:
+            self._rejected_nodes[node] = reason
+
+    def get_table_report(self) -> str:
+        """Returns a string containing a table listing all rejected nodes.
+        The table looks something like this:
+        ╒══════════════════════════╤══════════════════════════╤═════════════════════════════════════╤═════════════════════════════════════╕
+        │ Node name                │ Target                   │ Torch func                          │ Reason                              │
+        ╞══════════════════════════╪══════════════════════════╪═════════════════════════════════════╪═════════════════════════════════════╡
+        │ aten_convolution_default │ aten.convolution.default │ ('conv2d_1', 'builtin_function_or_m │ Convolution needs to have           │
+        │                          │                          │ ethod.conv2d')                      │ kernel_y<=64,                       │
+        │                          │                          │                                     │ kernel_x*kernel_y<=4096, got kernel │
+        │                          │                          │                                     │ (2, 65)                             │
+        ╘══════════════════════════╧══════════════════════════╧═════════════════════════════════════╧═════════════════════════════════════╛
+        """
+        reject_report = []
+        for node in self._rejected_nodes:
+            if node.op == "placeholder" or node.op == "output":
+                continue
+            if not (target := getattr(node.target, "_op", None)):
+                target = node.target
+            torch_fn = node.meta.get("torch_fn", "-")
+            reject_report.append(
+                [node.name, target, torch_fn, self._rejected_nodes[node]]
+            )
+        if len(reject_report) > 0:
+            return tabulate(
+                reject_report,
+                ["Node name", "Target", "Torch func", "Reason"],
+                tablefmt="fancy_grid",
+                maxcolwidths=35,
+            )
+        else:
+            return "No nodes rejected."
+
+    def wrap_check(
+        self, operator_support: OperatorSupportBase, message: str
+    ) -> OperatorSupportBase:
+        """Wrap the operator_support, reporting rejects with the specified message."""
+        return ReportRejected(operator_support, self, message)
+
+
+class ReportRejected(OperatorSupportBase):
+    """Class for wrapping a OperatorSupportBase, reporting rejects with the specified message to `reporter`."""
+
+    def __init__(
+        self,
+        operator_support: OperatorSupportBase,
+        reporter: WhyNoPartitionReporter,
+        message,
+    ):
+        self.operator_support = operator_support
+        self.reporter = reporter
+        self.message = message
+
+    def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
+        is_supported = self.operator_support.is_node_supported(submodules, node)
+        if not is_supported:
+            self.reporter.report_reject(node, self.message)
+        return is_supported

executorch/exir/capture/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+
+# pyre-strict
+
+from executorch.exir.capture._capture import (
+    _capture_legacy_do_not_use,
+    CallSpec,
+    capture,
+)
+
+from executorch.exir.capture._config import (
+    CaptureConfig,
+    EdgeCompileConfig,
+    ExecutorchBackendConfig,
+)
+
+__all__ = [
+    "CallSpec",
+    "capture",
+    "_capture_legacy_do_not_use",
+    "CaptureConfig",
+    "EdgeCompileConfig",
+    "ExecutorchBackendConfig",
+]