NXP backend: Do not infer format for unknown nodes.

Author: MartinPavella

backends/nxp/backend/edge_program_converter.py

@@ -19,7 +19,7 @@
 from torch.nn.parameter import Parameter
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters import *  # noqa F403
 from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
-from executorch.backends.nxp.backend.node_format_inference import NXP_NODE_FORMAT
+from executorch.backends.nxp.backend.node_format import NXP_NODE_FORMAT
 from executorch.exir.dialects._ops import ops as exir_ops
 
 # noinspection PyProtectedMember

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

@@ -9,7 +9,7 @@
 from executorch.backends.nxp.backend.ir.converter.conversion import translator
 from executorch.backends.nxp.backend.ir.tensor_formatting import TensorFormat
 from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model
-from executorch.backends.nxp.backend.node_format_inference import NodeFormat
+from executorch.backends.nxp.backend.node_format import NodeFormat
 from torch.fx import Node
 from torch.nn import Parameter
 

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

@@ -21,7 +21,7 @@
     Concatenation,
 )
 from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
-from executorch.backends.nxp.backend.node_format_inference import NXP_NODE_FORMAT
+from executorch.backends.nxp.backend.node_format import NXP_NODE_FORMAT
 from torch.fx import Node
 from torch.nn import Parameter
 

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

@@ -28,7 +28,7 @@
 )
 from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 
-from executorch.backends.nxp.backend.node_format_inference import NXP_NODE_FORMAT
+from executorch.backends.nxp.backend.node_format import NXP_NODE_FORMAT
 from torch.fx import Node
 from torch.nn import Parameter
 

backends/nxp/backend/ir/tensor_formatting.py

@@ -6,7 +6,7 @@
 #
 from enum import Enum
 
-from executorch.backends.nxp.backend.node_format_inference import NodeFormat
+from executorch.backends.nxp.backend.node_format import NodeFormat
 
 
 class TensorFormat(Enum):

backends/nxp/backend/node_format.py

@@ -0,0 +1,23 @@
+# 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.
+
+from enum import Enum
+
+# Key into the `meta` attribute of nodes, which is mapped to their inferred node format.
+NXP_NODE_FORMAT = "nxp_node_format"
+
+
+class NodeFormat(Enum):
+    # Node's output in NCHW format
+    CHANNELS_FIRST = 0
+
+    # Node's output format has no meaning
+    FORMATLESS = 1
+
+    # Format has not been identified
+    NONE = 2
+
+    def is_channels_first(self) -> bool:
+        return self == NodeFormat.CHANNELS_FIRST

backends/nxp/backend/node_format_inference.py

@@ -4,31 +4,18 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
-from enum import Enum
+import operator
 
+from executorch.backends.nxp.backend.edge_program_converter import functions_converters
+from executorch.backends.nxp.backend.node_format import NodeFormat, NXP_NODE_FORMAT
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
 
-from torch import Node
 from torch.export import ExportedProgram
+from torch.fx import Node
 
 logger = logging.getLogger(__name__)
 
-NXP_NODE_FORMAT = "nxp_node_format"  # Key into the `meta` attribute of nodes, which is mapped to the inferred format.
-
-
-class NodeFormat(Enum):
-    # Node's output in NCHW format
-    CHANNELS_FIRST = 0
-
-    # Node's output format has no meaning
-    FORMATLESS = 1
-
-    # Format has not been identified
-    NONE = 2
-
-    def is_channels_first(self) -> bool:
-        return self == NodeFormat.CHANNELS_FIRST
-
 
 class NodeFormatInference:
     # Dictionary with Edge Aten ops that always use channels first format.
@@ -53,6 +40,9 @@ class NodeFormatInference:
     # Mapping between Node and its children (outputs)
     _node_outputs: dict[Node, list[Node]]
 
+    # List of all edge operations, which are supported by the converter.
+    _known_targets: list[EdgeOpOverload]
+
     def __init__(self, edge_program: ExportedProgram):
         self._edge_program = edge_program
 
@@ -66,6 +56,13 @@ def __init__(self, edge_program: ExportedProgram):
 
         self._type_changed_during_last_run = False
 
+        self._known_targets = list(functions_converters) + [
+            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+            exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
+            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+            operator.getitem,
+        ]
+
     def identify_node_formats(self):
         self._type_changed_during_last_run = True
 
@@ -100,9 +97,19 @@ def _infer_format_of_nodes(self, node: Node):
                 logger.error(
                     f"Node format inference for node type: {op_type} not found!"
                 )
-        else:
+        elif node.op != "call_function" or (
+            hasattr(node, "target") and node.target in self._known_targets
+        ):
+            # Generic node, or tensor.
             self._handle_node_which_can_use_any_node_format(node)
 
+        else:
+            # Don't infer the format for unknown nodes. These nodes will never be delegated, so they will divide
+            #  delegated partitions. Propagating the format here could unnecessarily enforce the format in one of these
+            #  partitions, which would require extra transpositions.
+            for processed_node in self._node_inputs[node] + [node]:
+                self._assign_format_to_node(processed_node, NodeFormat.NONE)
+
     def _infer_format_based_on_io_ranks(self, node: Node):
         """Determine the format of the output tensor of given "reshape style operator" based on the ranks of its input
         and output.
@@ -155,6 +162,10 @@ def _assign_format_to_node(self, node: Node, node_format: NodeFormat):
             # Once CHANNEL_FIRST was assigned, we don't want to reassign
             return
 
+        if node_format is NodeFormat.NONE and old_node_format is not NodeFormat.NONE:
+            # A format has already been assigned to the node before. Don't replace it with `NONE`.
+            return
+
         if old_node_format != node_format:
             self._type_changed_during_last_run = True