NXP backend: Store inferred node format in the node.meta.

Author: MartinPavella

backends/nxp/backend/edge_program_converter.py

@@ -19,8 +19,8 @@
 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_inference import (
-    NodeFormat,
     NodeFormatInference,
+    NXP_NODE_FORMAT,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 
@@ -70,12 +70,11 @@ 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()
+        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,
         )
@@ -101,7 +100,7 @@ def convert_program(
     def append_placeholders_and_tensors(nodes: list[Node], context: ConversionContext):
         for node in nodes:
             if node.op == "placeholder":
-                node_format = context.node_formats[node]
+                node_format = node.meta[NXP_NODE_FORMAT]
 
                 if node.name in context.parameters_mapping:
                     # Node is placeholder and has data -> append as static tensor with data
@@ -114,7 +113,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 = context.node_formats[node]
+                node_format = node.meta[NXP_NODE_FORMAT]
                 context.tflite_builder.append_as_fake_tensor(node, node_format)
             elif node.op == "output":
                 # Nothing to do
@@ -172,7 +171,6 @@ 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:
@@ -187,7 +185,6 @@ def build_conversion_context(
             tflite_builder,
             conversion_config,
             parameters_mapping,
-            node_formats,
             custom_delegation_options,
         )
 

backends/nxp/backend/ir/conversion_context.py

@@ -10,24 +10,20 @@
 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,
     ):
         """
@@ -39,5 +35,4 @@ 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/node_format_inference.py

@@ -13,6 +13,8 @@
 
 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
@@ -43,8 +45,6 @@ class NodeFormatInference:
     # are channels first but output is formatless).
     ops_that_can_change_tensor_format = {exir_ops.edge.aten.view_copy.default}
 
-    _node_format_mapping: dict[Node, NodeFormat]
-
     _type_changed_during_last_run: bool
 
     # Mapping between Node and its ancestors (inputs)
@@ -57,7 +57,6 @@ def __init__(self, edge_program: ExportedProgram):
         self._edge_program = edge_program
 
         self._nodes = edge_program.graph.nodes
-        self._node_format_mapping = {}
         self._node_inputs = {
             node: node.all_input_nodes for node in edge_program.graph.nodes
         }
@@ -67,7 +66,7 @@ def __init__(self, edge_program: ExportedProgram):
 
         self._type_changed_during_last_run = False
 
-    def identify_node_formats(self) -> dict[Node, NodeFormat]:
+    def identify_node_formats(self):
         self._type_changed_during_last_run = True
 
         # Re-run format inference until there are no changes
@@ -77,7 +76,15 @@ def identify_node_formats(self) -> dict[Node, NodeFormat]:
             for node in self._nodes:
                 self._infer_format_of_nodes(node)
 
-        return self._node_format_mapping
+        for node in self._nodes:
+            if self._get_node_op_type(node) is None:
+                continue
+            if not hasattr(node, "meta"):
+                logging.warning(f"Node `{node}` does not have the `meta` attribute.")
+                node.meta = {}
+            if NXP_NODE_FORMAT not in node.meta:
+                logging.warning(f"Node `{node}` does not have inferred format.")
+                node.meta[NXP_NODE_FORMAT] = NodeFormat.NONE
 
     def _infer_format_of_nodes(self, node: Node):
         op_type = self._get_node_op_type(node)
@@ -151,7 +158,7 @@ def _assign_format_to_node(self, node: Node, node_format: NodeFormat):
         if old_node_format != node_format:
             self._type_changed_during_last_run = True
 
-        self._node_format_mapping[node] = node_format
+        node.meta[NXP_NODE_FORMAT] = node_format
 
     def _get_node_op_type(self, node: Node) -> str | None:
         """
@@ -252,8 +259,10 @@ def _node_produces_or_consumes_channels_first_format(self, node) -> bool:
             for ancestor_node in input_nodes
         )
 
-    def _get_node_format(self, node):
-        return self._node_format_mapping.get(node, NodeFormat.NONE)
+    def _get_node_format(self, node) -> NodeFormat:
+        if not hasattr(node, "meta"):
+            node.meta = {}
+        return node.meta.get(NXP_NODE_FORMAT, NodeFormat.NONE)
 
-    def _node_is_placeholder(self, node: Node):
+    def _node_is_placeholder(self, node: Node) -> bool:
         return node.op == "placeholder"

backends/nxp/tests/test_node_format_inference.py

@@ -9,6 +9,7 @@
 from executorch.backends.nxp.backend.node_format_inference import (
     NodeFormat,
     NodeFormatInference,
+    NXP_NODE_FORMAT,
 )
 from executorch.backends.nxp.neutron_pass_manager import NeutronPassManager
 from executorch.backends.nxp.tests.models import (
@@ -27,7 +28,7 @@ def test_convolution():
     exir_program = torch.export.export(model, example_input)
     edge_program = exir.to_edge(exir_program).exported_program()
 
-    node_formats = NodeFormatInference(edge_program).identify_node_formats()
+    NodeFormatInference(edge_program).identify_node_formats()
 
     expected_mapping = {
         "p_conv_weight": NodeFormat.CHANNELS_FIRST,
@@ -37,8 +38,8 @@ def test_convolution():
         "output": NodeFormat.CHANNELS_FIRST,
     }
 
-    for node, node_format in node_formats.items():
-        assert expected_mapping[node.name] == node_format
+    for node in edge_program.graph.nodes:
+        assert expected_mapping[node.name] == node.meta[NXP_NODE_FORMAT]
 
 
 def test_softmax():
@@ -48,16 +49,16 @@ def test_softmax():
     exir_program = torch.export.export(model, example_input)
     edge_program = exir.to_edge(exir_program).exported_program()
 
-    node_formats = NodeFormatInference(edge_program).identify_node_formats()
+    NodeFormatInference(edge_program).identify_node_formats()
 
     expected_mapping = {
         "x": NodeFormat.FORMATLESS,
         "aten__softmax_default": NodeFormat.FORMATLESS,
         "output": NodeFormat.FORMATLESS,
     }
 
-    for node, node_format in node_formats.items():
-        assert expected_mapping[node.name] == node_format
+    for node in edge_program.graph.nodes:
+        assert expected_mapping[node.name] == node.meta[NXP_NODE_FORMAT]
 
 
 def test_maxpool2d():
@@ -78,13 +79,13 @@ def test_maxpool2d():
 
     # Remove MaxPool-related "getitem" nodes from graph
     edge_program = NeutronPassManager(edge_program, [RemoveGetItemPass]).transform()
-    node_formats = NodeFormatInference(edge_program).identify_node_formats()
+    NodeFormatInference(edge_program).identify_node_formats()
 
     expected_mapping = {
         "x": NodeFormat.CHANNELS_FIRST,
         "aten_max_pool2d_default": NodeFormat.CHANNELS_FIRST,
         "output": NodeFormat.CHANNELS_FIRST,
     }
 
-    for node, node_format in node_formats.items():
-        assert expected_mapping[node.name] == node_format
+    for node in edge_program.graph.nodes:
+        assert expected_mapping[node.name] == node.meta[NXP_NODE_FORMAT]