NXP backend: Add DetailedGraphVerifier for strict checking of delegated nodes.

Author: MartinPavella

backends/nxp/tests/graph_verifier.py

@@ -5,42 +5,85 @@
 
 import abc
 import re
+from collections import defaultdict
+from copy import deepcopy
 from dataclasses import dataclass
-from typing import Union
+from typing import Callable, Union
 
+from executorch.backends.nxp.neutron_partitioner import (
+    NeutronPartitioner,
+    NXP_DELEGATION_TAG,
+)
+from executorch.backends.nxp.tests.ops_aliases import (
+    DequantizePerChannel,
+    DequantizePerTensor,
+    QuantizePerChannel,
+    QuantizePerTensor,
+)
+
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
+
+from pytest_mock import MockerFixture
+
+from torch.fx import Node
 from torch.fx.graph import Graph
 
 
 @dataclass
 class NonDelegatedNode:
+    """Represents an expected non-delegated node in the graph.
+
+    :param node_name: The name of the node to check for
+    :param num_occurrences: Expected number of occurrences. If None, just verifies that at least one exists
+    """
+
     node_name: str
     num_occurrences: Union[int, None] = None
 
 
 class GraphVerifier(abc.ABC):
+    """Abstract base class for graph verification strategies."""
+
     @abc.abstractmethod
     def verify_graph(self, graph: Graph):
-        pass
+        """Verifies the graph meets expected criteria.
 
-    @abc.abstractmethod
-    def check_num_delegated_nodes(self, num_dlg_nodes: int):
+        :param graph: The FX graph to verify
+        :raises AssertionError: If the graph does not meet expectations
+        """
         pass
 
 
 class BaseGraphVerifier(GraphVerifier):
-    """Graph verifier base class. Checks for number of delegated nodes and number of selected expected nodes."""
+    """Graph verifier base class. Checks for number of delegated nodes and number of selected expected nodes.
+
+    This verifier performs the following checks:
+    - The total number of delegated call nodes matches expectations
+    - Specific non-delegated nodes appear with the expected frequency
+    - No unexpected aten nodes are present in the graph
+    """
 
     def __init__(
         self,
         exp_num_delegate_call_nodes: int,
         exp_non_delegated_nodes: list[NonDelegatedNode] = None,
     ):
+        """Initializes the BaseGraphVerifier.
+
+        :param exp_num_delegate_call_nodes: Expected number of delegated nodes
+        :param exp_non_delegated_nodes: List of expected non-delegated nodes to verify
+        """
         self.exp_non_delegated_nodes = (
             exp_non_delegated_nodes if exp_non_delegated_nodes is not None else []
         )
         self.exp_num_delegate_call_nodes = exp_num_delegate_call_nodes
 
     def check_num_delegated_nodes(self, num_dlg_nodes):
+        """Checks that the number of delegated nodes matches expectations.
+
+        :param num_dlg_nodes: Actual number of delegated nodes
+        :raises AssertionError: If the count doesn't match expectations
+        """
         assert not (
             num_dlg_nodes < self.exp_num_delegate_call_nodes
         ), f"Number of delegated nodes decreased from {self.exp_num_delegate_call_nodes} to {num_dlg_nodes}."
@@ -49,6 +92,11 @@ def check_num_delegated_nodes(self, num_dlg_nodes):
         ), f"Number of delegated nodes increased from {self.exp_num_delegate_call_nodes} to {num_dlg_nodes}."
 
     def verify_graph(self, graph):
+        """Verifies the graph meets delegation and node presence expectations.
+
+        :param graph: The FX graph to verify
+        :raises AssertionError: If verification fails
+        """
         nodes = list(graph.nodes)
 
         # Check for specific non delegated nodes
@@ -84,3 +132,132 @@ def verify_graph(self, graph):
         assert (
             not unexpected_aten_fn_nodes
         ), f"Graphs contains unexpected aten nodes:\n{unexpected_aten_fn_nodes}."
+
+
+# Type alias for operators - can be either EdgeOpOverload or any callable (e.g., operator.getitem).
+Operator = EdgeOpOverload | Callable
+
+
+class DetailedGraphVerifier(GraphVerifier):
+    """Graph verifier that checks for exact delegated and non-delegated operators.
+
+    This verifier captures a snapshot of the graph immediately after partitioning and verifies
+    that specific operators were delegated/non-delegated the expected number of times. It uses
+    mocker to intercept the partition() call and create a deep copy of the nodes before they
+    can be modified. Quantization/dequantization operators are ignored by default as they are
+    typically not the focus of delegation verification.
+    """
+
+    default_ops_to_ignore = {
+        QuantizePerTensor,
+        QuantizePerChannel,
+        DequantizePerTensor,
+        DequantizePerChannel,
+    }
+
+    def __init__(
+        self,
+        expected_delegated_ops: dict[Operator, int],
+        expected_non_delegated_ops: dict[Operator, int],
+        mocker: MockerFixture,
+        ops_to_ignore: set[Operator] | None = None,
+    ):
+        """Initializes the DetailedGraphVerifier and patches NeutronPartitioner.partition() to capture node state.
+
+        :param expected_delegated_ops: Dictionary mapping operators to their expected delegation count
+        :param expected_non_delegated_ops: Dictionary mapping operators to their expected non-delegation count
+        :param mocker: Pytest mocker fixture for intercepting the partition method
+        :param ops_to_ignore: Set of operators to ignore during verification. Defaults to quantization ops
+        """
+        self.expected_delegated_ops = expected_delegated_ops
+        self.expected_non_delegated_ops = expected_non_delegated_ops
+
+        self.ops_to_ignore = ops_to_ignore or self.default_ops_to_ignore
+
+        # We need to use mocker to capture a copy of the nodes returned by NeutronPartitioner.partition() to access
+        # their partition tag. The nodes in the returned graph may be modified after partition() returns, so we
+        # capture a deep copy immediately when the method completes.
+        self.captured_partitioned_nodes: list[Node] | None = None
+
+        # Store original partition method for the wrapper.
+        # Note: pytest-mock automatically restores the original method after the test completes,
+        # so manual cleanup is not required.
+        original_partition_method = NeutronPartitioner.partition
+
+        def partition_wrapper(self_, exported_program):
+            """Wraps NeutronPartitioner.partition() to capture a snapshot of nodes after partitioning.
+
+            :param self_: The NeutronPartitioner instance
+            :param exported_program: The ExportedProgram being partitioned
+            :return: The PartitionResult from the original partition method
+            """
+            result = original_partition_method(self_, exported_program)
+            # Capture a deep copy of the nodes with their metadata.
+            # This ensures we have the exact state immediately after partitioning,
+            # before any subsequent transformations modify the graph.
+            self.captured_partitioned_nodes = list(
+                deepcopy(exported_program.graph.nodes)
+            )
+            return result
+
+        # Patch the partition method to intercept and capture results.
+        mocker.patch.object(NeutronPartitioner, "partition", partition_wrapper)
+
+    def verify_graph(self, graph):
+        """Verifies that operators were delegated/non-delegated as expected by comparing actual counts against expectations.
+
+        :param graph: The FX graph to verify (not directly used; we use captured nodes instead)
+        :raises AssertionError: If the NeutronPartitioner wasn't used or if delegation doesn't match expectations
+        """
+        assert (
+            self.captured_partitioned_nodes is not None
+        ), "The NeutronPartitioner was not used. Cannot access delegated nodes."
+
+        delegated_ops = defaultdict(int)
+        non_delegated_ops = defaultdict(int)
+
+        for node in self.captured_partitioned_nodes:
+            # Only process call_function nodes with a target
+            if not hasattr(node, "target") or node.op != "call_function":
+                continue
+
+            # Skip operators we're configured to ignore (e.g., quantization ops)
+            if node.target in self.ops_to_ignore:
+                continue
+
+            # Check if the node was tagged for delegation during partitioning
+            if NXP_DELEGATION_TAG in node.meta:
+                delegated_ops[node.target] += 1
+            else:
+                non_delegated_ops[node.target] += 1
+
+        # All ops which were either expected to be delegated, or were actually delegated.
+        all_delegated_ops = list(set(self.expected_delegated_ops).union(delegated_ops))
+
+        # All ops which were either expected to be non-delegated, or were actually non-delegated.
+        all_non_delegated_ops = list(
+            set(self.expected_non_delegated_ops).union(non_delegated_ops)
+        )
+
+        message = ""
+
+        # Check delegated operators
+        for op in all_delegated_ops:
+            expected_count = self.expected_delegated_ops.get(op, 0)
+            real_count = delegated_ops.get(op, 0)
+            op_name = op.name() if hasattr(op, "name") else str(op)
+            if expected_count != real_count:
+                message += f"\t`{op_name}` was delegated {real_count} times instead of the expected {expected_count} times.\n"
+
+        # Check non-delegated operators
+        for op in all_non_delegated_ops:
+            expected_count = self.expected_non_delegated_ops.get(op, 0)
+            real_count = non_delegated_ops.get(op, 0)
+            op_name = op.name() if hasattr(op, "name") else str(op)
+            if expected_count != real_count:
+                message += f"\t`{op_name}` was NON-delegated {real_count} times instead of the expected {expected_count} times.\n"
+
+        if message:
+            raise AssertionError(
+                "Some operators were not delegated as expected:\n" + message
+            )

backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py

@@ -28,10 +28,7 @@
     ToNCHWPreprocess,
     ToNHWCPreprocess,
 )
-from executorch.backends.nxp.tests.graph_verifier import (
-    BaseGraphVerifier,
-    NonDelegatedNode,
-)
+from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier
 from executorch.backends.nxp.tests.models import AvgPool2dConvModule, AvgPool2dModule
 
 from executorch.backends.nxp.tests.nsys_testing import lower_run_compare
@@ -306,26 +303,20 @@ def test_from_avg_pool_1d(mocker):
 
 
 class TestAvgPool2DNewNeutronFlow:
-    def test__basic_nsys_inference(self):
+    def test__basic_nsys_inference(self, mocker):
         input_shape = (2, 4, 6, 7)
         model = AvgPool2dModule(False, 0)
-        graph_verifier = BaseGraphVerifier(
-            exp_num_delegate_call_nodes=1,  # Delegated AvgPool.
-            exp_non_delegated_nodes=[],
-        )
+        graph_verifier = DetailedGraphVerifier({AvgPool2D: 1}, {}, mocker)
 
         lower_run_compare(
             model, input_shape, graph_verifier, use_new_flow_neutron_c=True
         )
 
-    def test__kernel_size_limit(self):
+    def test__kernel_size_limit(self, mocker):
         kernel_size = (1, 4096)
         input_shape = (1, 4) + kernel_size
         model = AvgPool2dModule(False, 0, kernel_size)
-        graph_verifier = BaseGraphVerifier(
-            exp_num_delegate_call_nodes=1,  # Delegated AvgPool.
-            exp_non_delegated_nodes=[],
-        )
+        graph_verifier = DetailedGraphVerifier({AvgPool2D: 1}, {}, mocker)
 
         lower_run_compare(
             model, input_shape, graph_verifier, use_new_flow_neutron_c=True
@@ -346,14 +337,11 @@ def test__kernel_size_limit_exceeded(self):
         )
         assert graph_contains_any_of_ops(delegated_ep.graph, [AvgPool2D])
 
-    def test__stride_limit(self):
+    def test__stride_limit(self, mocker):
         stride = 4096
         input_shape = (1, 4, 1, 4096)
         model = AvgPool2dModule(False, 0, 1, stride)
-        graph_verifier = BaseGraphVerifier(
-            exp_num_delegate_call_nodes=1,  # Delegated AvgPool.
-            exp_non_delegated_nodes=[],
-        )
+        graph_verifier = DetailedGraphVerifier({AvgPool2D: 1}, {}, mocker)
 
         lower_run_compare(
             model, input_shape, graph_verifier, use_new_flow_neutron_c=True
@@ -378,17 +366,10 @@ def test__stride_limit_exceeded(self):
 class TestAvgPool1DNewNeutronFlow:
 
     # Just a basic test to verify that the operator gets extended to the 2D variant correctly.
-    def test__basic_nsys_inference__view_not_delegated(self):
+    def test__basic_nsys_inference__view_not_delegated(self, mocker):
         input_shape = (2, 4, 6)  # The old flow limited the batch size to 1.
         model = AvgPool1DModule()
-        graph_verifier = BaseGraphVerifier(
-            exp_num_delegate_call_nodes=1,  # Delegated AvgPool.
-            exp_non_delegated_nodes=[
-                NonDelegatedNode(
-                    "aten_view_copy_default", 2
-                )  # Non delegated due to shape requirements.
-            ],
-        )
+        graph_verifier = DetailedGraphVerifier({AvgPool2D: 1}, {ViewCopy: 2}, mocker)
 
         lower_run_compare(
             model, input_shape, graph_verifier, use_new_flow_neutron_c=True