NXP backend: Add infrastructure for context dependant partitioning

backends/nxp/backend/ir/converter/node_converter.py

@@ -18,6 +18,7 @@
 from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx import Node
+from torch.fx.passes.infra.partitioner import Partition
 from torch.nn import Parameter
 
 
@@ -37,6 +38,10 @@ def _is_dequant_node(node: torch.fx.Node) -> bool:
     ]
 
 
+def is_not_qdq_node(node: torch.fx.Node) -> bool:
+    return not (_is_quant_node(node) or _is_dequant_node(node))
+
+
 class Target(Enum):
     IGNORE = "ignore"  # No target platform. Any target specific restrictions will be ignored.
 
@@ -125,6 +130,23 @@ def is_supported(
             node, target, parameters_mapping, custom_delegation_options
         )
 
+    @classmethod
+    def supports_partitioning_result(
+        cls,
+        node: Node,
+        partition_list: list[Partition],
+        custom_delegation_options: CustomDelegationOptions,
+    ):
+        """Check if the given `node` supports the assigned partitioning, which is stored  the `partition_list`. Child
+            classes can overwrite this method in case they have delegation restrictions based on the context defined by
+            the partitioning result.
+
+        :param node: torch.Node to check.
+        :param partition_list: List of proposed partitions.
+        :param custom_delegation_options: Custom user options which affect node delegation.
+        """
+        return True
+
     @staticmethod
     def _has_shared_q_params_if_quantized(node: Node) -> bool:
         """Check if node has shared quantization parameters if it's quantized."""

backends/nxp/backend/ir/converter/node_converters/ops_converters/view_copy_converter.py

@@ -14,6 +14,7 @@
 from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
+    is_not_qdq_node,
     NodeConverter,
 )
 from executorch.backends.nxp.backend.ir.converter.node_converters.shared.reshape_transposition import (
@@ -23,6 +24,7 @@
     reshape_options,
 )
 from torch.fx import Node
+from torch.fx.passes.infra.partitioner import Partition
 from torch.nn import Parameter
 
 
@@ -45,6 +47,27 @@ def _is_supported_in_IR(
 
         return True
 
+    @classmethod
+    def supports_partitioning_result(
+        cls,
+        node: Node,
+        partition_list: list[Partition],
+        custom_delegation_options: CustomDelegationOptions,
+    ):
+        view_copy_partitions = [
+            partition for partition in partition_list if node in partition.nodes
+        ]
+        assert len(view_copy_partitions) == 1
+        non_q_dq_partition_nodes = list(
+            filter(is_not_qdq_node, view_copy_partitions[0].nodes)
+        )
+
+        if len(non_q_dq_partition_nodes) == 1:
+            # The `view_copy` cannot be the only node in a partition.
+            return False
+
+        return True
+
     @staticmethod
     def _safe_compute_flat_size(shape: list[int | str]) -> int:
         """Compute the flat size of a tensor with given shape. Strings and negative dimensions are treated as '1'.

backends/nxp/neutron_partitioner.py

@@ -20,7 +20,8 @@
 )
 from executorch.backends.nxp.backend.ir.converter.node_converter import Target
 from torch.export.exported_program import ExportedProgram
-from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
+from torch.fx import Graph
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
 from torch.nn import Parameter
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters import *  # noqa F403
@@ -34,6 +35,9 @@
 from executorch.exir.backend.utils import tag_constant_data
 from executorch.exir.dialects._ops import ops as exir_ops
 
+NXP_DO_NOT_DELEGATE = "NXP_DO_NOT_DELEGATE"
+NXP_DELEGATION_TAG = "delegation_tag"
+
 
 class QDQClusterRecognizer:
     """
@@ -246,6 +250,11 @@ def _is_node_supported_compute(self, node: torch.fx.node.Node) -> bool:
         """
         Operator checking function for compute nodes.
         """
+
+        if hasattr(node, "meta") and node.meta.get(NXP_DO_NOT_DELEGATE, False):
+            # The delegation of this node has been prohibited.
+            return False
+
         if not self.is_node_delegatable(node):
             return False
 
@@ -304,6 +313,31 @@ def __init__(
             custom_delegation_options or CustomDelegationOptions()
         )
 
+    def validate_partitioning_result(
+        self,
+        graph: Graph,
+        partition_list: list[Partition],
+        custom_delegation_options: CustomDelegationOptions,
+    ) -> bool:
+        all_delegated_nodes = {
+            node for partition in partition_list for node in partition.nodes
+        }
+        partitioning_valid = True
+        for node in graph.nodes:
+            if (
+                node in all_delegated_nodes
+                and hasattr(node, "target")
+                and node.target in supported_ops
+            ):
+                if not supported_ops[node.target].supports_partitioning_result(
+                    node, partition_list, custom_delegation_options
+                ):
+                    # This node is not supported within its partition. Exclude it from delegation in the future.
+                    partitioning_valid = False
+                    node.meta[NXP_DO_NOT_DELEGATE] = True
+
+        return partitioning_valid
+
     def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         # Run the CapabilityBasedPartitioner to return the largest possible
         # subgraphs containing the nodes with the tags
@@ -342,11 +376,24 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
             allows_single_node_partition=True,
         )
 
-        partition_list = capability_partitioner.propose_partitions()
+        iteration_limit = len(exported_program.graph.nodes)
+        for _ in range(iteration_limit):
+            # Run the partitioning.
+            partition_list = capability_partitioner.propose_partitions()
+
+            # Check if the nodes support the partitioning result. Mark the problematic nodes with `NXP_DO_NOT_DELEGATE`.
+            partitioning_valid = self.validate_partitioning_result(
+                exported_program.graph, partition_list, self.custom_delegation_options
+            )
+            if partitioning_valid:
+                # The result of the partitioning is fine
+                break
+
+        # Mark the partitions in the node `meta` attribute.
         for partition in partition_list:
             for node in partition.nodes:
                 delegation_tag = f"tag{partition.id}"
-                node.meta["delegation_tag"] = delegation_tag
+                node.meta[NXP_DELEGATION_TAG] = delegation_tag
                 partition_tags[delegation_tag] = self.delegation_spec
 
         tag_constant_data(exported_program)

backends/nxp/tests/test_context_sensitive_delegation.py

@@ -0,0 +1,71 @@
+# Copyright 2025 NXP
+#
+# 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
+
+import numpy as np
+import torch
+
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters import (
+    ViewCopyConverter,
+)
+from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
+from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+class SingleViewCopyModule(torch.nn.Module):
+    def __init__(self, new_shape: list[int]):
+        super().__init__()
+        self.new_shape = new_shape
+
+    def forward(self, x):
+        return torch.reshape(x, self.new_shape)
+
+
+class TestContextSensitiveDelegation(unittest.TestCase):
+    __test__ = False  # Prevent interfering with PyTest tests.
+
+    @classmethod
+    def setUpClass(cls):
+        torch.manual_seed(23)
+        np.random.seed(42)
+
+    def test_single_view_copy_partition(self):
+        input_shape = (2, 10)
+        module = SingleViewCopyModule([1, 20])
+
+        ep = to_quantized_edge_program(module, input_shape).exported_program()
+
+        # Make sure the `view_copy` was not delegated.
+        assert graph_contains_any_of_ops(
+            ep.graph, [exir_ops.edge.aten.view_copy.default]
+        )
+        assert not any("delegate" in n.name for n in ep.graph.nodes)
+
+    def test_single_view_copy_partition__forced_delegation(self):
+        input_shape = (2, 10)
+        module = SingleViewCopyModule([1, 20])
+
+        def _supported_partitioning(*_):
+            return True
+
+        # Replace the partition support check function, to accept anything.
+        original_supports_partitioning_result = (
+            ViewCopyConverter.supports_partitioning_result
+        )
+        ViewCopyConverter.supports_partitioning_result = _supported_partitioning
+
+        with self.assertRaises(RuntimeError) as e:
+            to_quantized_edge_program(module, input_shape).exported_program()
+        assert (
+            str(e.exception)
+            == "Model converted with neutron-converter does not contain a NeutronGraph node."
+        )
+
+        # Return to the original partition support check function.
+        ViewCopyConverter.supports_partitioning_result = (
+            original_supports_partitioning_result
+        )