diff --git a/backends/nxp/backend/ir/converter/node_converter.py b/backends/nxp/backend/ir/converter/node_converter.py index f8405f37680..a1e12ec08eb 100755 --- a/backends/nxp/backend/ir/converter/node_converter.py +++ b/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 diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py index 471fb7a1f22..ef6d66504bf 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py +++ b/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] diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py index ea3914f4fe2..888a2e6e689 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py +++ b/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 diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py index 5fa994be7ae..b02583f2f62 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py +++ b/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( diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py index a30475d64c3..39384078df4 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py +++ b/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 diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py index d57124247b4..d7be6b0efa8 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py +++ b/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 diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py index 64d0601824c..e361e38e1c8 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py +++ b/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 diff --git a/backends/nxp/tests/generic_tests/test_node_format_inference.py b/backends/nxp/tests/generic_tests/test_node_format_inference.py index f588e464fa9..ecdff416a09 100644 --- a/backends/nxp/tests/generic_tests/test_node_format_inference.py +++ b/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