NXP backend: Fix bug when a partition output is also used by nodes inside the partition.

Author: MartinPavella

backends/nxp/backend/edge_program_converter.py

@@ -106,6 +106,9 @@ def convert_program(
         self._convert_qdq_cluster_q_dq_nodes(edge_program.graph.nodes, cc)
         self._process_nodes(edge_program.graph.nodes, cc)
 
+        # Make sure the operators are correctly ordered.
+        cc.tflite_builder.sort_operators_topologically()
+
         # Assign the model its inputs and outputs.
         cc.tflite_builder.assign_model_io_to_subgraph(edge_program.graph_signature)
 

backends/nxp/backend/ir/converter/builder/model_builder.py

@@ -6,6 +6,8 @@
 # See the LICENSE_MIT for more details.
 #
 
+import logging
+from collections import defaultdict, deque
 from copy import deepcopy
 from typing import List, Optional, Union
 
@@ -106,6 +108,65 @@ def __init__(
         self._skipped_output_map = {}
         self._zeros_tensor_map = {}
 
+    def sort_operators_topologically(self):
+        """Sort the operators in the model graph in topological order using Kahn's algorithm.
+
+        Ensures that each operator appears after all operators that produce its input tensors.
+        This is required for correct behavior of optimization passes and by the Neutron IR format.
+
+        The algorithm works in three steps:
+            1. Build a mapping from each tensor to the operator that produces it.
+            2. Compute the number of unresolved input dependencies for each operator.
+            3. Iteratively add operators with no remaining dependencies to the sorted list,
+               decrementing the number of dependencies of their consumers until all operators are processed.
+
+        If the sort succeeds, the operator list in the model's subgraph is replaced with the topologically sorted list.
+        If it fails (e.g. due to an invalid graph), a warning is logged and the original operator order is preserved.
+        """
+
+        # Map tensors to operators that produce them.
+        tensor_to_producer_op = {}
+        for producer_op in self.get_operators():
+            for output in producer_op.tmp_outputs:
+                tensor_to_producer_op[output] = producer_op
+
+        # Map operators to operators that consume their outputs.
+        op_to_child_op = defaultdict(list)
+        num_unresolved_input_dependencies = {op: 0 for op in self.get_operators()}
+        for child_op in self.get_operators():
+            for input_ in child_op.tmp_inputs:
+                parent_op = tensor_to_producer_op.get(input_)
+                if parent_op is not None:
+                    op_to_child_op[parent_op].append(child_op)
+                    num_unresolved_input_dependencies[child_op] += 1
+
+        # Gradually build the graph.
+        # Use `deque` for fast appends and pops.
+        queue = deque(
+            op
+            for op, degree in num_unresolved_input_dependencies.items()
+            if degree == 0
+        )
+        sorted_ops = []
+        while len(queue) != 0:
+            op = queue.popleft()
+            sorted_ops.append(op)
+            for child_op in op_to_child_op[op]:
+                num_unresolved_input_dependencies[child_op] -= 1
+                if num_unresolved_input_dependencies[child_op] == 0:
+                    # All input of the `child_op` are produced by operators that are already in `sorted_ops`. So we can
+                    #  insert it into the graph.
+                    queue.append(child_op)
+
+        if len(sorted_ops) != self.get_operators().len():
+            logging.warning(
+                "NXP backend: ModelBuilder.sort_operators_topologically() failed. Please report this."
+            )
+
+        else:
+            # The topological sort was successful.
+            self.get_operators().vector = sorted_ops
+
     def create_zeros_tensor(
         self, dims: List[int], name: str, dtype: np.dtype, can_reuse: bool = False
     ) -> tflite_model.Tensor:

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

@@ -60,12 +60,13 @@ def convert(self, node: Node):
         if isinstance(self, QDQPerChannelDequantizeConverter):
             quantized_dimension = self.get_quantization_dimension(input_tensor, node)
 
-        consumes_model_input = (
+        consumes_model_io = (
             node.args[0].name in self.context.edge_program_signature.user_inputs
+            or node.args[0].name in self.context.edge_program_signature.user_outputs
         )
-        if consumes_model_input:
-            # We cannot just skip the operator. Skipping would require changing the input's name, and as the input is
-            #  also a model input, the name cannot be changed.
+        if consumes_model_io:
+            # We cannot just skip the operator. Skipping would require changing the input's/output's name, and as the
+            #  input/output is also a model input/output, the name cannot be changed.
             # Instead, we convert it into an identity (Transpose that will be removed), and we make the output tensor
             #  quantized just like the input.
             t_op = self._create_tflite_op_with_io_tensors(node)
@@ -80,6 +81,7 @@ def convert(self, node: Node):
 
             self.builder.turn_operator_to_identity(t_op)
             self.builder.append_operators([t_op])
+
         else:
             # Dequantize consumes an internal tensor, so we can just make it so that any operators which used the float
             #  output of the dequantize will now use its quantized input. We do this by redirecting the output to the

backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py

@@ -8,7 +8,6 @@
 import executorch.extension.pybindings.portable_lib
 import executorch.kernels.quantized  # noqa F401
 
-import pytest
 import torch
 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
 from executorch.backends.nxp.tests.models import Conv2dReLUModule
@@ -95,7 +94,6 @@ def forward(self, x, y):
         return x + y, z
 
 
-@pytest.mark.xfail(strict=True, reason="Known bug (EIEX-946).")
 def test_multiple_inputs__multiple_outputs():
     model = MultiInputOutputModule()
     model.eval()

backends/nxp/tests/ir/test_model_builder.py

@@ -0,0 +1,69 @@
+# Copyright 2026 NXP
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from copy import deepcopy
+
+from executorch.backends.nxp.backend.ir.converter.builder.model_builder import (
+    ModelBuilder,
+)
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import (
+    Transpose,
+)
+from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
+from executorch.backends.nxp.tests.ir.edge_passes.test_remove_io_quant_ops_pass import (
+    MultiInputOutputModule,
+)
+
+
+def test_topological_sorting(mocker):
+    # This model requires the insertion of a `Transpose` operator in place of a `dequantize` node, which breaks the
+    #  topological order.
+    model = MultiInputOutputModule()
+
+    # Capture the operators before and after topological sorting.
+    original_sort_fn = ModelBuilder.sort_operators_topologically
+    captured_ops = {}
+
+    def mock_sort_fn(self: ModelBuilder):
+        captured_ops["pre_sort"] = deepcopy(self.get_operators().vector)
+        original_sort_fn(self)  # type: ignore
+        captured_ops["post_sort"] = deepcopy(self.get_operators().vector)
+
+    mocker.patch.object(
+        ModelBuilder,
+        "sort_operators_topologically",
+        autospec=True,
+        side_effect=mock_sort_fn,
+    )
+
+    input_shapes = [(1, 4, 32, 32), (1, 1, 1, 31)]
+    to_quantized_edge_program(model, input_shapes)
+
+    ops_pre_sort = captured_ops["pre_sort"]
+    ops_post_sort = captured_ops["post_sort"]
+    assert len(ops_pre_sort) == len(ops_post_sort)
+
+    # Before sorting, the operator on index `2` should be an identity `Transpose`, which uses the output of the
+    #  operator on index `3` (breaking the topological order).
+    assert isinstance(ops_pre_sort[2].builtin_options, Transpose)
+    assert all(ops_pre_sort[2].tmp_inputs[1].tmp_buffer.data == [0, 1, 2, 3])
+    assert ops_pre_sort[2].tmp_inputs[0] == ops_pre_sort[3].tmp_outputs[0]
+
+    # After the sort, the operators on indices `2` and `3` are swapped.
+    assert (
+        ops_post_sort[2].builtin_options is None
+    )  # A Relu, which doesn't have a `BuiltinOptions` in Neutron IR.
+    assert isinstance(ops_post_sort[3].builtin_options, Transpose)
+    assert all(ops_post_sort[3].tmp_inputs[1].tmp_buffer.data == [0, 1, 2, 3])
+    assert ops_post_sort[2].tmp_outputs[0] == ops_post_sort[3].tmp_inputs[0]
+
+    # Make sure all nodes follow topological order.
+    tensor_to_producer_op = {
+        output_tensor: op for op in ops_post_sort for output_tensor in op.tmp_outputs
+    }
+    for op_idx, op in enumerate(ops_post_sort):
+        for input_tensor in op.tmp_inputs:
+            if input_tensor in tensor_to_producer_op:
+                assert ops_post_sort.index(tensor_to_producer_op[input_tensor]) < op_idx