NXP backend: Add format checks to prevent NodeFormatInference bugs. (#20333)

### Summary
Add format checks to prevent NodeFormatInference bugs. In the past we
have had PRs which introduced new ops but failed to update the format
inference. The decorator introduced in this PR should reduce the
likelihood of it happening again.

### Test plan
Unit tests provided.


cc @robert-kalmar @JakeStevens @digantdesai @rascani

Author: MartinPavella

backends/nxp/backend/ir/converter/node_converter.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 import operator
 from abc import ABC, abstractmethod
 from typing import Callable
@@ -12,6 +13,7 @@
 from executorch.backends.nxp.backend.custom_delegation_options import (
     CustomDelegationOptions,
 )
+from executorch.backends.nxp.backend.data_format import DataFormat, NXP_NODE_FORMAT
 from executorch.backends.nxp.backend.edge_helper import (
     input_quantization_type,
     output_quantization_type,
@@ -53,6 +55,23 @@ def is_not_qdq_node(node: torch.fx.Node) -> bool:
     return not (_is_quant_node(node) or _is_dequant_node(node))
 
 
+def requires_channels_first_format(cls):
+    """Class decorator for NodeConverter subclasses.
+
+    Marks a converter as requiring that both the node's main input and output
+    use the channels-first data format (as inferred by NodeFormatInference).
+    The check is automatically enforced via `NodeConverter.is_supported()`.
+
+    Usage::
+
+        @requires_channels_first_format
+        class ConvConverter(NodeConverter):
+            ...
+    """
+    cls._requires_channels_first_format = True
+    return cls
+
+
 class NodeConverter(ABC):
     """
     Classes which implement conversion of torch.Node to TFLite should inherit from this class and overwrite the
@@ -61,6 +80,11 @@ class NodeConverter(ABC):
 
     context: ConversionContext
 
+    # If `True`, the `is_supported()` method will disallow delegation if the node's main input/output doesn't have the
+    #  channels first node format.
+    # Subclasses decorated with @requires_channels_first_format will have this set to True.
+    _requires_channels_first_format: bool = False
+
     def __init__(self, context: ConversionContext):
         self.context = context
 
@@ -115,6 +139,36 @@ def _is_supported_on_target(
         """
         return True
 
+    @classmethod
+    def _node_format_is_supported(cls, node: Node) -> bool:
+        """Check that the node's main input and output carry the channels-first data format, if the converter was
+             decorated with `@requires_channels_first_format`.
+
+        When the decorator is not present the check returns True.
+
+        :param node: The node to inspect.
+        :return: True when the format requirement is satisfied (or not applicable).
+        """
+        if not cls._requires_channels_first_format:
+            return True
+
+        def _is_channels_first(n: Node) -> bool:
+            return (
+                n.meta.get(NXP_NODE_FORMAT, DataFormat.NONE)
+                is DataFormat.CHANNELS_FIRST
+            )
+
+        format_requirement_satisfied = _is_channels_first(node) and _is_channels_first(
+            node.args[0]
+        )
+        if not format_requirement_satisfied:
+            logging.warning(
+                f"NXP backend: Node `{node}` requires channels-first format for its input and output, but the inferred "
+                "format does not satisfy this requirement. The node will not be delegated. Please report this issue."
+            )
+
+        return format_requirement_satisfied
+
     @classmethod
     def is_supported(
         cls,
@@ -133,10 +187,15 @@ def is_supported(
                                     be outdated.
         :param custom_delegation_options: Custom user options which affect node delegation.
         """
-        return cls._is_supported_in_IR(
-            node, parameters_mapping, custom_delegation_options
-        ) and cls._is_supported_on_target(
-            node, neutron_target_spec, parameters_mapping, custom_delegation_options
+
+        return (
+            cls._node_format_is_supported(node)
+            and cls._is_supported_in_IR(
+                node, parameters_mapping, custom_delegation_options
+            )
+            and cls._is_supported_on_target(
+                node, neutron_target_spec, parameters_mapping, custom_delegation_options
+            )
         )
 
     @classmethod
@@ -225,22 +284,16 @@ def assert_convertible(self, node):
         """Assert that the call `is_supported()` returns `True`. Otherwise, raise an exception and print an
         error message.
         """
-        supported_in_ir = self._is_supported_in_IR(
-            node,
-            self.context.parameters_mapping,
-            self.context.custom_delegation_options,
-        )
 
-        supported_on_target = self._is_supported_on_target(
+        is_supported = self.is_supported(
             node,
             self.neutron_target_spec,
             self.context.parameters_mapping,
             self.context.custom_delegation_options,
         )
-
-        assert supported_in_ir and supported_on_target, (
-            f"Node `{node}` was selected for delegation to Neutron, but it is not convertible to the intermediate "
-            "representation. There is an error in the Neutron partitioner. Please report this."
+        assert is_supported, (
+            f"NXP backend: Node `{node}` was selected for delegation to Neutron, but it is not convertible to the "
+            "intermediate representation. There is an error in the Neutron partitioner. Please report this."
         )
 
     @property

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

@@ -2,16 +2,15 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-import logging
 
 import executorch.backends.nxp.backend.ir.lib.tflite.Padding as tflPadding
 import torch
 
-from executorch.backends.nxp.backend.data_format import NXP_NODE_FORMAT
 from executorch.backends.nxp.backend.ir.converter.conversion import common
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     NodeConverter,
+    requires_channels_first_format,
 )
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
     average_pool_2d_options,
@@ -25,6 +24,7 @@
 Stride = tuple[int, int]
 
 
+@requires_channels_first_format
 class AdaptiveAvgPool2dConverter(NodeConverter):
 
     @staticmethod
@@ -45,15 +45,6 @@ def _is_supported_in_IR(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if (
-            format_ := node.meta.get(NXP_NODE_FORMAT)
-        ) is None or not format_.is_channels_first():
-            logging.warning(
-                "NXP backend: `adaptive_avg_pool_2d` doesn't have the required input format for delegation. "
-                "Please run `NodeFormatInference.identify_node_formats()` during lowering or report this issue."
-            )
-            return False
-
         input_size = node.args[0].meta["val"].shape
         output_size = node.args[1]
 

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

@@ -16,6 +16,7 @@
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     NodeConverter,
+    requires_channels_first_format,
 )
 from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
@@ -26,6 +27,7 @@
 from torch.nn import Parameter
 
 
+@requires_channels_first_format
 class AvgPool2dConverter(NodeConverter):
 
     @staticmethod

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

@@ -23,6 +23,7 @@
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     NodeConverter,
+    requires_channels_first_format,
 )
 from executorch.backends.nxp.backend.ir.converter.node_converters.shared import (
     conv_utils,
@@ -62,6 +63,7 @@
 ]
 
 
+@requires_channels_first_format
 class ConvolutionConverter(NodeConverter):
     @staticmethod
     def _is_supported_on_target_regular_conv(

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

@@ -16,6 +16,7 @@
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     NodeConverter,
+    requires_channels_first_format,
 )
 from executorch.backends.nxp.backend.ir.lib.tflite.TensorType import TensorType
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.max_pool_2d_options import (
@@ -32,6 +33,7 @@
 CeilMode = bool
 
 
+@requires_channels_first_format
 class MaxPool2DWithIndicesConverter(NodeConverter):
 
     @staticmethod

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

@@ -6,12 +6,12 @@
 import numpy as np
 import torch
 
-from executorch.backends.nxp.backend.data_format import DataFormat, NXP_NODE_FORMAT
 from executorch.backends.nxp.backend.edge_helper import node_has_well_defined_shape
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     is_not_qdq_node,
     NodeConverter,
+    requires_channels_first_format,
 )
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.resize_bilinear_options import (
     ResizeBilinear,
@@ -23,6 +23,7 @@
 
 
 # noinspection SpellCheckingInspection
+@requires_channels_first_format
 class UpsampleBilinear2DConverter(NodeConverter):
 
     @classmethod
@@ -53,14 +54,6 @@ def _is_supported_in_IR(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-
-        if node.meta.get(NXP_NODE_FORMAT, DataFormat.NONE) != DataFormat.CHANNELS_FIRST:
-            # This should never happen.
-            raise NotImplementedError(
-                "NXP backend: `aten.upsample_bilinear2d.vec` didn't have correctly identified data"
-                " format. Please report this."
-            )
-
         # The conversion requires the output shape to be known and static.
         if not node_has_well_defined_shape(node):
             return False

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

@@ -6,12 +6,12 @@
 import numpy as np
 import torch
 
-from executorch.backends.nxp.backend.data_format import DataFormat, NXP_NODE_FORMAT
 from executorch.backends.nxp.backend.edge_helper import node_has_well_defined_shape
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     is_not_qdq_node,
     NodeConverter,
+    requires_channels_first_format,
 )
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.resize_nearest_neighbor_options import (
     ResizeNearestNeighbor,
@@ -26,6 +26,7 @@
 
 
 # noinspection SpellCheckingInspection
+@requires_channels_first_format
 class UpsampleNearest2DConverter(NodeConverter):
 
     @classmethod
@@ -55,14 +56,6 @@ def _is_supported_in_IR(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-
-        if node.meta.get(NXP_NODE_FORMAT, DataFormat.NONE) != DataFormat.CHANNELS_FIRST:
-            # This should never happen.
-            raise NotImplementedError(
-                "NXP backend: `aten.upsample_nearest2d.vec` didn't have correctly identified data"
-                " format. Please report this."
-            )
-
         # The conversion requires the output shape to be known and static.
         if not node_has_well_defined_shape(node):
             return False

backends/nxp/tests/generic_tests/test_node_format_inference.py

@@ -3,6 +3,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
+
 import torch
 
 from executorch import exir
@@ -11,12 +13,18 @@
     NodeFormatInference,
     NXP_NODE_FORMAT,
 )
+from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
+from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops
 
 from executorch.backends.nxp.tests.models import (
     Conv2dModule,
     MaxPool2dModule,
     SoftmaxModule,
 )
+from executorch.backends.nxp.tests.ops_aliases import (
+    ExecutorchDelegateCall,
+    MaxPool2DWithIndices,
+)
 
 
 def test_convolution():
@@ -77,3 +85,37 @@ def test_max_pool2d():
 
     for node in epm.exported_program().graph.nodes:
         assert expected_mapping[node.name] == node.meta[NXP_NODE_FORMAT]
+
+
+def test_unhandled_channels_first_node(caplog):
+    # This test focuses on the case where some operator requires the channels first format, which is enforced in the
+    #  `NodeConverter`, but the `NodeFormatInference` fails to reflect this.
+    # We use the `MaxPool` operator for this, and we temporarily modify the `NodeFormatInference` to trigger the issue.
+
+    model = MaxPool2dModule()
+    input_shape = (1, 4, 32, 32)
+
+    # Temporarily "break" the NodeFormatInference.
+    old_channels_first_ops = NodeFormatInference.ops_with_channels_first_nodes
+    NodeFormatInference.ops_with_channels_first_nodes = {}
+
+    try:
+        with caplog.at_level(
+            logging.WARNING,
+            logger="executorch.backends.nxp.backend.ir.converter.node_converter",
+        ):
+            ep = to_quantized_edge_program(model, input_shape).exported_program()
+    finally:
+        # Restore the original channels first ops configuration.
+        NodeFormatInference.ops_with_channels_first_nodes = old_channels_first_ops
+
+    # Make sure the `MaxPool` wasn't delegated.
+    assert graph_contains_any_of_ops(ep.graph, [MaxPool2DWithIndices])
+    assert not graph_contains_any_of_ops(ep.graph, [ExecutorchDelegateCall])
+
+    # Make sure the warning is printed.
+    assert any(
+        "`aten_max_pool2d_with_indices_default` requires channels-first format for its input and output, but the "
+        "inferred format does not satisfy this requirement" in message
+        for message in caplog.messages
+    )