Revert "NXP backend: Resolve limitations of uncertain tensor formats"

backends/nxp/backend/edge_program_converter.py

@@ -18,7 +18,10 @@
 from torch.fx import Node
 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.node_format import NXP_NODE_FORMAT
+from executorch.backends.nxp.backend.node_format_inference import (
+    NodeFormat,
+    NodeFormatInference,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 
 # noinspection PyProtectedMember
@@ -67,10 +70,12 @@ def convert_program(
         :param custom_delegation_options: Custom user options which affect node delegation.
         :return: TFLite flatbuffers as bytes.
         """
+        node_formats = NodeFormatInference(edge_program).identify_node_formats()
         parameters_mapping = self.map_inputs_to_parameters(edge_program)
 
         cc = self.build_conversion_context(
             parameters_mapping,
+            node_formats,
             conversion_config,
             custom_delegation_options,
         )
@@ -96,7 +101,7 @@ def convert_program(
     def append_placeholders_and_tensors(nodes: list[Node], context: ConversionContext):
         for node in nodes:
             if node.op == "placeholder":
-                node_format = node.meta[NXP_NODE_FORMAT]
+                node_format = context.node_formats[node]
 
                 if node.name in context.parameters_mapping:
                     # Node is placeholder and has data -> append as static tensor with data
@@ -109,7 +114,7 @@ def append_placeholders_and_tensors(nodes: list[Node], context: ConversionContex
                     context.tflite_builder.append_as_fake_tensor(node, node_format)
             elif node.op == "call_function":
                 # Node is call function -> append only output as a tensor
-                node_format = node.meta[NXP_NODE_FORMAT]
+                node_format = context.node_formats[node]
                 context.tflite_builder.append_as_fake_tensor(node, node_format)
             elif node.op == "output":
                 # Nothing to do
@@ -166,6 +171,7 @@ def map_inputs_to_parameters(edge_program: ExportedProgram) -> dict[str, Paramet
     @staticmethod
     def build_conversion_context(
         parameters_mapping: dict,
+        node_formats: dict[Node, NodeFormat],
         conversion_config: ConversionConfig = _default_conversion_config,
         custom_delegation_options: CustomDelegationOptions = _default_delegation_options,
     ) -> ConversionContext:
@@ -180,6 +186,7 @@ def build_conversion_context(
             tflite_builder,
             conversion_config,
             parameters_mapping,
+            node_formats,
             custom_delegation_options,
         )
 

backends/nxp/backend/ir/conversion_context.py

@@ -10,20 +10,24 @@
 from executorch.backends.nxp.backend.ir.converter.builder.aten_model_builder_director import (
     AtenModelBuilderDirector,
 )
+from executorch.backends.nxp.backend.node_format_inference import NodeFormat
+from torch import Node
 from torch.nn import Parameter
 
 
 class ConversionContext:
     tflite_builder: AtenModelBuilderDirector
     conversion_config: ConversionConfig
     parameters_mapping: dict[str, Parameter]
+    node_formats: dict[Node, NodeFormat]
     custom_delegation_options: CustomDelegationOptions
 
     def __init__(
         self,
         tflite_builder: AtenModelBuilderDirector,
         conversion_config: ConversionConfig,
         parameters_mapping: dict,
+        node_formats: dict[Node, NodeFormat],
         custom_delegation_options: CustomDelegationOptions,
     ):
         """
@@ -35,4 +39,5 @@ def __init__(
         self.tflite_builder = tflite_builder
         self.conversion_config = conversion_config
         self.parameters_mapping = parameters_mapping
+        self.node_formats = node_formats
         self.custom_delegation_options = custom_delegation_options

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 import NodeFormat
+from executorch.backends.nxp.backend.node_format_inference import NodeFormat
 from torch.fx import Node
 from torch.nn import Parameter
 

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

@@ -18,7 +18,6 @@
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.concatenation_options import (
     Concatenation,
 )
-from executorch.backends.nxp.backend.node_format import NXP_NODE_FORMAT
 from torch.fx import Node
 from torch.nn import Parameter
 
@@ -89,27 +88,25 @@ def _is_supported_on_target(
                     return False
 
                 # Neutron requires the channels to be a multiple of `8`. The channels could either be the second or the
-                #  last dimension, depending on the formats of the node.
-                if node.meta[NXP_NODE_FORMAT].is_channels_first():
-                    # During conversion to IR, the shape will be permuted to channels last, and the dimension on index
-                    #  `1` will end up being the channels (last dim in NHWC).
-                    channels_index = 1
-                else:
-                    # The shape will not be permuted during conversion, so the channels will remain the last dimension.
-                    channels_index = -1
-
+                #  last dimension, depending on the formats of the node. The format, however, cannot be determined
+                #  during conversion, as it depends on what other nodes are delegated.
                 input_channels = [
-                    _get_shape(input_)[channels_index]
+                    # The second dimension is the channels in PyTorch. If the inputs/output are not channels first, it
+                    #  will still be the channels in the IR.
+                    _get_shape(input_)[1]
+                    for input_ in node.all_input_nodes
+                ] + [
+                    # If the inputs/outputs are channels first, the last dimension will be the channels.
+                    _get_shape(input_)[-1]
                     for input_ in node.all_input_nodes
                 ]
-                output_channels = _get_shape(node)[channels_index]
-
                 if any((input_channel % 8) != 0 for input_channel in input_channels):
                     # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1492
                     return False
 
-                if (output_channels % 8) != 0:
-                    # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1493
+                output_channels = [_get_shape(node)[1], _get_shape(node)[-1]]
+                # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1493
+                if any((out_c % 8) != 0 for out_c in output_channels):
                     return False
 
                 if len(node.all_input_nodes) < 2:  # Not supported on Neutron

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

@@ -27,8 +27,6 @@
     pad_options,
     pad_v2_options,
 )
-
-from executorch.backends.nxp.backend.node_format import NXP_NODE_FORMAT
 from torch.fx import Node
 from torch.nn import Parameter
 
@@ -43,17 +41,11 @@ def _is_supported_on_target(
     ) -> bool:
         match target:
             case Target.RT700:
+                # TODO: Consider different tensor formats (dim-order)
                 paddings = node.args[1]
-                if node.meta[NXP_NODE_FORMAT].is_channels_first():
-                    # Dim `1` will end up being the channels. It is padded by paddings[4:6].
-                    if len(paddings) > 4 and paddings[4:6] != [0, 0]:
-                        # Attempt to Pad channels dimension -> currently not supported
-                        return False
-                else:
-                    # Dim `-1` will end up being the channels. It is padded by paddings[:2].
-                    if len(paddings) > 0 and paddings[:2] != [0, 0]:
-                        # Attempt to Pad channels dimension -> currently not supported
-                        return False
+                if len(paddings) > 4 and paddings[4:6] != [0, 0]:
+                    # Attempt to Pad channels dimension, which is not supported on Neutron.
+                    return False
 
                 return True
 
@@ -79,6 +71,10 @@ def _is_supported_in_IR(
         if not NodeConverter._has_shared_q_params_if_quantized(node):
             return False
 
+        if len(paddings) > 4 and paddings[4:6] != [0, 0]:
+            # Attempt to Pad channels dimension -> currently not supported
+            return False
+
         return True
 
     # noinspection PyMethodMayBeStatic

backends/nxp/backend/ir/tensor_formatting.py

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