NXP backend: Perform node format inference before partitioning.

Before, the format inference was done during conversion to NeutronIR (after partitioning), so the partitioner didn't yet know the formats. Now, the partitioner has the format data, which can be used to accurately select nodes for delegation.

Author: MartinPavella

backends/nxp/backend/edge_program_converter.py

@@ -19,10 +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 (
-    NodeFormatInference,
-    NXP_NODE_FORMAT,
-)
+from executorch.backends.nxp.backend.node_format_inference import NXP_NODE_FORMAT
 from executorch.exir.dialects._ops import ops as exir_ops
 
 # noinspection PyProtectedMember
@@ -74,7 +71,6 @@ def convert_program(
         :param custom_delegation_options: Custom user options which affect node delegation.
         :return: TFLite flatbuffers as bytes.
         """
-        NodeFormatInference(edge_program).identify_node_formats()
         parameters_mapping = self.map_inputs_to_parameters(edge_program)
 
         cc = self.build_conversion_context(

backends/nxp/neutron_partitioner.py

@@ -25,6 +25,7 @@
 from torch.nn 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 NodeFormatInference
 from executorch.backends.nxp.nxp_backend import NeutronBackend
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.backend.partitioner import (
@@ -374,6 +375,10 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
             allows_single_node_partition=True,
         )
 
+        # Identify the format (NCHW/NHWC/...) for all nodes in the graph, and store it in the `node.meta`.
+        # This format will be used by the `CapabilityBasedPartitioner` to determine which nodes will be delegated.
+        NodeFormatInference(exported_program).identify_node_formats()
+
         iteration_limit = len(exported_program.graph.nodes)
         for _ in range(iteration_limit):
             # Run the partitioning.

backends/nxp/tests/executors.py

@@ -24,7 +24,6 @@
 from torch.fx import Node
 from torch.fx.graph import Graph
 
-
 # If executed on i.MX platform, there is no tensorflow module. And typically the intention is to use the tflite python
 # interpreter available in tflite_runtime
 try:
@@ -308,6 +307,7 @@ def convert_run_compare(
 ) -> (TFLiteExecutor, EdgeProgramExecutor):
 
     if tfl_model is None:
+        NodeFormatInference(edge_program).identify_node_formats()
         tfl_model, _ = EdgeProgramToIRConverter().convert_program(
             edge_program, conversion_config
         )

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

@@ -17,6 +17,8 @@
 from executorch.backends.nxp.tests.executors import (
     convert_run_compare,
     graph_contains_any_of_ops,
+    ToNCHWPreprocess,
+    ToNHWCPreprocess,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import ExportedProgram
@@ -126,6 +128,8 @@ def test_cat__channels_first__same_shapes(dim, num_inputs, mocker):
         exported_program,
         tfl_model=tflite_flatbuffers_model,
         input_data=input_data,
+        tflite_input_preprocess=ToNHWCPreprocess(),
+        tflite_output_preprocess=ToNCHWPreprocess(),
         atol=1,
     )
 
@@ -241,6 +245,8 @@ def test_cat__channels_first__different_shapes(dim, num_inputs, mocker):
         exported_program,
         tfl_model=tflite_flatbuffers_model,
         input_data=input_data,
+        tflite_input_preprocess=ToNHWCPreprocess(),
+        tflite_output_preprocess=ToNCHWPreprocess(),
         atol=1,
     )
 

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

@@ -1,4 +1,4 @@
-# Copyright 2024 NXP
+# Copyright 2024-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.
@@ -11,6 +11,7 @@
     EdgeProgramToIRConverter,
 )
 from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig
+from executorch.backends.nxp.backend.node_format_inference import NodeFormatInference
 from executorch.backends.nxp.tests.executorch_pipeline import to_edge_program
 from executorch.backends.nxp.tests.executors import convert_run_compare
 from executorch.backends.nxp.tests.models import SoftmaxConvModule, SoftmaxModule
@@ -56,6 +57,7 @@ def test_softmax_conversion__unknown_input_format(input_shape, dim: int):
     model = SoftmaxModule(dim)
 
     edge_program = to_edge_program(model, input_shape).exported_program()
+    NodeFormatInference(edge_program).identify_node_formats()
 
     # Currently this test not pass because the convertibility checker doesn't use tensor formats.
     with pytest.raises(
@@ -78,6 +80,7 @@ def test_softmax_conversion_channel_last(input_shape, dim: int):
     model = SoftmaxConvModule(dim)
 
     edge_program = to_edge_program(model, input_shape).exported_program()
+    NodeFormatInference(edge_program).identify_node_formats()
 
     # TODO (Robert Kalmar) Currently this test not pass because the convertibility checker doesn't use tensor formats.
     with pytest.raises(
@@ -104,6 +107,7 @@ def test_softmax_conversion_unsupported_dims(input_shape, dim: int):
     model = SoftmaxModule(dim)
 
     edge_program = to_edge_program(model, input_shape).exported_program()
+    NodeFormatInference(edge_program).identify_node_formats()
 
     with pytest.raises(
         AssertionError, match="`aten__softmax_default` is not convertible"

backends/nxp/tests/test_neutron_converter_manager.py

@@ -13,6 +13,7 @@
 from executorch.backends.nxp.backend.neutron_converter_manager import (
     NeutronConverterManager,
 )
+from executorch.backends.nxp.backend.node_format_inference import NodeFormatInference
 from executorch.backends.nxp.tests.models import Conv2dModule
 
 
@@ -23,6 +24,7 @@ def test_conv2d_neutron_conversion__default_flavor():
     exir_program = torch.export.export(model, example_input)
     edge_program_manager = exir.to_edge(exir_program)
 
+    NodeFormatInference(edge_program_manager.exported_program()).identify_node_formats()
     edge_program_converter = EdgeProgramToIRConverter()
     tflite_model, _ = edge_program_converter.convert_program(
         edge_program_manager.exported_program()
@@ -43,6 +45,7 @@ def test__conv2d_neutron_conversion__invalid_flavor():
     exir_program = torch.export.export(model, example_input)
     edge_program_manager = exir.to_edge(exir_program)
 
+    NodeFormatInference(edge_program_manager.exported_program()).identify_node_formats()
     edge_program_converter = EdgeProgramToIRConverter()
     tflite_model, _ = edge_program_converter.convert_program(
         edge_program_manager.exported_program()