NXP backend: Add format checks to prevent NodeFormatInference bugs.

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,13 @@ 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._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
+            )
+            and cls._node_format_is_supported(node)
         )
 
     @classmethod

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,
@@ -48,6 +49,7 @@
 from torch.nn import Parameter
 
 
+@requires_channels_first_format
 class ConvolutionConverter(NodeConverter):
     @staticmethod
     def _is_supported_on_target(

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,36 @@ 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 = {}
+
+    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()
+
+    # 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
+    )
+
+    # Restore the original channels first ops configuration.
+    NodeFormatInference.ops_with_channels_first_nodes = old_channels_first_ops