Pull request pytorch#124: Implementation of 'aten::hardtanh' operator conversion

Merge in AITEC/executorch from feature/nxg10272/EIEX-239-add-conversion-for-hardtanh-aka-relu6 to main-nxp

* commit '3bc4775a02cd1c47e322402d7fd6eca495e47ec6':
  Add tests, implementation, and integration of 'aten::hardtanh' operator conversion

Author: StrycekSimon

backends/nxp/backend/edge_program_converter.py

@@ -29,6 +29,7 @@
     exir_ops.edge.aten.mm.default: MMConverter,
     exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter,
     exir_ops.edge.aten.relu.default: ReLUConverter,
+    exir_ops.edge.aten.hardtanh.default: HardTanhConverter,
     exir_ops.edge.aten._softmax.default: SoftmaxConverter,
     exir_ops.edge.aten.view_copy.default: ViewCopyConverter,
     exir_ops.edge.aten.add.Tensor: AddTensorConverter,

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

@@ -30,9 +30,11 @@
     CloneConverter
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.abs_converter import \
     AbsConverter
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.hardtanh_converter import \
+    HardTanhConverter
 __all__ = [
     "AddMMConverter", "ConvolutionConverter", "MMConverter", "PermuteCopyConverter", "SoftmaxConverter",
     "ViewCopyConverter", "QDQDequantizeConverter", "QDQQuantizeConverter", "ConstantPadNDConverter", "ReLUConverter",
     "MaxPool2dConverter", "AvgPool2dConverter", "AddTensorConverter", "MeanDimConverter", "AdaptiveAvgPool2dConverter",
-    "CloneConverter", "AbsConverter"
+    "CloneConverter", "AbsConverter", "HardTanhConverter"
 ]

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

@@ -0,0 +1,41 @@
+# Copyright (c) 2025 NXP
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch.fx import Node
+from torch.nn import Parameter
+
+from executorch.backends.nxp.backend.ir.converter.node_converter import NodeConverter, Target
+from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import BuiltinOperator
+
+
+class HardTanhConverter(NodeConverter):
+    supported_targets = [Target.RT700]
+
+    # Maps possible input parameters of HardTanh to equivalent ReLU-based operators supported by TFLite.
+    supported_modes_map = {
+        ( 0., 6.): BuiltinOperator.RELU6,
+        (-1., 1.): BuiltinOperator.RELU_N1_TO_1,
+        ( 0., 1.): BuiltinOperator.RELU_0_TO_1,
+        ( 0., float('inf')): BuiltinOperator.RELU,
+    }
+
+    @staticmethod
+    def _is_supported_in_IR(node: Node, parameters_mapping: dict[str, Parameter]) -> bool:
+        _, min_value, max_value = node.args
+        return (min_value, max_value) in HardTanhConverter.supported_modes_map.keys()
+
+    def convert(self, node: Node):
+        """ Convert 'aten::hardtanh' to it's supported ReLU equivalent. """
+        self.assert_convertible(node)
+
+        t_op = self._create_tflite_op_with_io_tensors(node)
+
+        _, min_value, max_value = node.args
+
+        op = self.supported_modes_map[(min_value, max_value)]
+        t_op.opcode_index = self.builder.op_code_index_for_op_type(op)
+
+        self.builder.append_operators([t_op])

backends/nxp/neutron_partitioner.py

@@ -187,6 +187,7 @@ def tag_qdq_clusters(self, nodes: List[torch.fx.Node]):
     exir_ops.edge.aten.max_pool2d_with_indices.default: MaxPool2dConverter,
     exir_ops.edge.aten.mm.default: MMConverter,
     exir_ops.edge.aten.relu.default: ReLUConverter,
+    exir_ops.edge.aten.hardtanh.default: HardTanhConverter,
     exir_ops.edge.aten._softmax.default: SoftmaxConverter,
     exir_ops.edge.aten.view_copy.default: ViewCopyConverter,
     exir_ops.edge.aten.add.Tensor: AddTensorConverter,

backends/nxp/quantizer/neutron_quantizer.py

@@ -169,6 +169,26 @@ def partition_types(self):
         return [torch.ops.aten.relu_.default]
 
 
+class HardTanhPattern(SharedSpecPattern):
+    """
+    Quantizer for HardTanh operator. Shared quantization spec is selected, as activation functions usually follows
+    computation layer.
+    """
+
+    def partition_types(self):
+        return [torch.ops.aten.hardtanh.default]
+
+
+class HardTanhInPlacePattern(SharedSpecPattern):
+    """
+    Quantizer for HardTanh operator with param inplace=True. Shared quantization spec is selected, as activation
+    functions usually follows computation layer.
+    """
+
+    def partition_types(self):
+        return [torch.ops.aten.hardtanh_.default]
+
+
 class ReshapePattern(SharedSpecPattern):
     """
     Quantizer for Reshape operator.
@@ -317,6 +337,8 @@ def __init__(self):
                 CadenceAtenQuantizer(PermutePattern(), static_qconfig),
                 CadenceAtenQuantizer(PadPattern(), static_qconfig),
                 CadenceAtenQuantizer(ReluPattern(), static_qconfig),
+                CadenceAtenQuantizer(HardTanhPattern(), static_qconfig),
+                CadenceAtenQuantizer(HardTanhInPlacePattern(), static_qconfig),
                 CadenceAtenQuantizer(ReluInPlacePattern(), static_qconfig),
                 CadenceAtenQuantizer(AvgPoolPattern(), static_qconfig),
                 CadenceAtenQuantizer(ViewPattern(), static_qconfig),

backends/nxp/tests/executors.py

@@ -267,8 +267,8 @@ def convert_run_compare(edge_program: ExportedProgram, input_data, rtol=1.e-5, a
     return tflite_executor, edge_program_executor
 
 
-def graph_contains_op(graph: Graph, op: object) -> bool:
-        return any(map(lambda node: node.target == op, graph.nodes))
+def graph_contains_any_of_ops(graph: Graph, ops: list) -> bool:
+        return any(map(lambda node: node.target in ops, graph.nodes))
 
 
 class OverrideSupportedTargets:

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

@@ -5,7 +5,7 @@
 
 from executorch.backends.nxp.backend.edge_program_converter import EdgeProgramToIRConverter
 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
-from executorch.backends.nxp.tests.executors import convert_run_compare, graph_contains_op, ToChannelLastPreprocess, \
+from executorch.backends.nxp.tests.executors import convert_run_compare, graph_contains_any_of_ops, ToChannelLastPreprocess, \
     ToChannelFirstPreprocess
 from executorch.exir.dialects._ops import ops as exir_ops
 
@@ -51,7 +51,7 @@ def test_conv_abs(mocker, input_shape: tuple[int] = (1, 3, 112, 112)):
     tflite_flatbuffers_model, io_formats = converter_spy.spy_return
     exported_program: ExportedProgram = converter_spy.call_args.args[1]
 
-    assert not graph_contains_op(graph=quantized_program.graph, op=exir_ops.edge.aten.abs.default)
+    assert not graph_contains_any_of_ops(graph=quantized_program.graph, ops=[exir_ops.edge.aten.abs.default])
 
     input_data = (np.random.random(input_shape) * 50).astype(np.int8)
     convert_run_compare(exported_program,
@@ -72,7 +72,7 @@ def test_abs_only(mocker, input_shape: tuple[int] = (1, 10)):
     tflite_flatbuffers_model, io_formats = converter_spy.spy_return
     exported_program: ExportedProgram = converter_spy.call_args.args[1]
 
-    assert not graph_contains_op(graph=quantized_program.graph, op=exir_ops.edge.aten.abs.default)
+    assert not graph_contains_any_of_ops(graph=quantized_program.graph, ops=[exir_ops.edge.aten.abs.default])
 
     input_data = (np.random.random(input_shape) * 50).astype(np.int8)
     convert_run_compare(exported_program,

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

@@ -12,7 +12,7 @@
 
 from executorch.backends.nxp.backend.edge_program_converter import EdgeProgramToIRConverter
 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
-from executorch.backends.nxp.tests.executors import convert_run_compare, graph_contains_op, ToChannelLastPreprocess, \
+from executorch.backends.nxp.tests.executors import convert_run_compare, graph_contains_any_of_ops, ToChannelLastPreprocess, \
     ToChannelFirstPreprocess
 from executorch.exir.dialects._ops import ops as exir_ops
 
@@ -75,7 +75,7 @@ def test_conv_dropout_quant(mocker, inplace_dropout: bool, input_shape: tuple[in
     tflite_flatbuffers_model, io_formats = converter_spy.spy_return
     exported_program: ExportedProgram = converter_spy.call_args.args[1]
 
-    assert not graph_contains_op(graph=quantized_program.graph, op=exir_ops.edge.aten.clone.default)
+    assert not graph_contains_any_of_ops(graph=quantized_program.graph, ops=[exir_ops.edge.aten.clone.default])
 
     input_data = (np.random.random(input_shape) * 50).astype(np.int8)
     convert_run_compare(exported_program,
@@ -97,7 +97,7 @@ def test_clone_pool_view_copy_quant(mocker, inplace_dropout: bool, input_shape:
     tflite_flatbuffers_model, io_formats = converter_spy.spy_return
     exported_program: ExportedProgram = converter_spy.call_args.args[1]
 
-    assert not graph_contains_op(graph=quantized_program.graph, op=exir_ops.edge.aten.clone.default)
+    assert not graph_contains_any_of_ops(graph=quantized_program.graph, ops=[exir_ops.edge.aten.clone.default])
 
     input_data = (np.random.random(input_shape) * 50).astype(np.int8)
     convert_run_compare(exported_program,

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

@@ -0,0 +1,101 @@
+import numpy as np
+import pytest
+import torch
+from torch.export import ExportedProgram
+
+from executorch.backends.nxp.backend.edge_program_converter import EdgeProgramToIRConverter
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.hardtanh_converter import \
+    HardTanhConverter
+from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
+from executorch.backends.nxp.tests.executors import convert_run_compare, ToChannelLastPreprocess, \
+    ToChannelFirstPreprocess, graph_contains_any_of_ops
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+@pytest.fixture(autouse=True)
+def reseed_model_per_test_run():
+    torch.manual_seed(23)
+    np.random.seed(23)
+
+
+class Relu6ConvBlock(torch.nn.Module):
+    def __init__(self, conv_in_channels: int = 3, inplace: bool = False):
+        super().__init__()
+        self.block = torch.nn.Sequential(
+          torch.nn.Conv2d(in_channels=conv_in_channels, out_channels=64, kernel_size=(4, 4)),
+          torch.nn.ReLU6(inplace=inplace)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class CustomHardTanhBlock(torch.nn.Module):
+    def __init__(self,
+                 conv_in_channels: int = 3,
+                 min_act_val: float = -1.,
+                 max_act_val: float = 1.,
+                 inplace: bool = False):
+        super().__init__()
+        self.block = torch.nn.Sequential(
+            torch.nn.Conv2d(in_channels=conv_in_channels, out_channels=64, kernel_size=(4, 4)),
+            torch.nn.Hardtanh(min_val=min_act_val, max_val=max_act_val, inplace=inplace)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+@pytest.mark.parametrize('input_shape', [(1, 3, 128, 128), (1, 3, 256, 256)])
+@pytest.mark.parametrize('inplace', [True, False])
+def test_relu6_quant(mocker, input_shape: tuple[int], inplace: bool):
+    model = Relu6ConvBlock(conv_in_channels=input_shape[1], inplace=inplace)
+
+    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+
+    quantized_program = to_quantized_edge_program(model, input_shape).exported_program()
+
+    tflite_flatbuffers_model, io_formats = converter_spy.spy_return
+    exported_program: ExportedProgram = converter_spy.call_args.args[1]
+
+    ops = [exir_ops.edge.aten.hardtanh.default, exir_ops.edge.aten.hardtanh_.default]
+    assert not graph_contains_any_of_ops(graph=quantized_program.graph, ops=ops)
+
+    input_data = (np.random.random(input_shape) * 50).astype(np.int8)
+    convert_run_compare(exported_program,
+                        tfl_model=tflite_flatbuffers_model,
+                        tflite_input_preprocess=ToChannelLastPreprocess(),
+                        tflite_output_preprocess=ToChannelFirstPreprocess(),
+                        input_data=input_data,
+                        atol=1.)
+
+
+@pytest.mark.parametrize('input_shape', [(1, 3, 128, 128), (1, 3, 256, 256)])
+@pytest.mark.parametrize('activation_range', list(HardTanhConverter.supported_modes_map.keys()))
+@pytest.mark.parametrize('inplace', [True, False])
+def test_custom_hardtanh_quant(mocker, input_shape: tuple[int], activation_range: tuple[int, int], inplace: bool):
+    min_val, max_val = activation_range
+    model = CustomHardTanhBlock(
+        conv_in_channels=input_shape[1],
+        min_act_val=min_val,
+        max_act_val=max_val,
+        inplace=inplace
+    )
+
+    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+
+    quantized_program = to_quantized_edge_program(model, input_shape).exported_program()
+
+    tflite_flatbuffers_model, io_formats = converter_spy.spy_return
+    exported_program: ExportedProgram = converter_spy.call_args.args[1]
+
+    ops = [exir_ops.edge.aten.hardtanh.default, exir_ops.edge.aten.hardtanh_.default]
+    assert not graph_contains_any_of_ops(graph=quantized_program.graph, ops=ops)
+
+    input_data = (np.random.random(input_shape) * 50).astype(np.int8)
+    convert_run_compare(exported_program,
+                        tfl_model=tflite_flatbuffers_model,
+                        tflite_input_preprocess=ToChannelLastPreprocess(),
+                        tflite_output_preprocess=ToChannelFirstPreprocess(),
+                        input_data=input_data,
+                        atol=1.)