Draft: Add extended support from new Neutron C flow for Clamp operator

backends/nxp/backend/custom_delegation_options.py

@@ -22,7 +22,3 @@ class CustomDelegationOptions:
     #  not create any NeutronGraph that can be called. This is done by the partitioner itself, and is not handled by
     #  the individual node converters.
     allow_no_op_partitions: bool = False
-
-    # The new neutron converter flow has different constraints for supported operators. These need to be addressed when
-    # deciding is operator is delegated or not in _is_supported_on_target().
-    use_new_flow_neutron_c: bool = False

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

@@ -5,7 +5,6 @@
 
 
 import torch
-
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     NeutronTargetSpec,
@@ -36,7 +35,7 @@ def _is_supported_on_target(
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
 
-        if custom_delegation_options.use_new_flow_neutron_c:
+        if neutron_target_spec.use_new_flow_neutron_c:
             # Requirements specified by the new Neutron flow documentation.
 
             supported_types = [torch.int8, torch.uint8]

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

@@ -5,7 +5,6 @@
 
 import numpy as np
 import torch
-
 from executorch.backends.nxp.backend.ir.converter.conversion import (
     aten_translator,
     common,
@@ -22,7 +21,6 @@
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
     average_pool_2d_options,
 )
-
 from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 from torch.fx import Node
 from torch.nn import Parameter
@@ -66,7 +64,7 @@ def _is_supported_on_target(
         kernel = node.args[1]
         stride = node.args[2]
 
-        if custom_delegation_options.use_new_flow_neutron_c:
+        if neutron_target_spec.use_new_flow_neutron_c:
             # Requirements specified by the new Neutron flow documentation.
 
             supported_types = [torch.int8, torch.uint8]

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

@@ -3,15 +3,24 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from copy import copy
+
+import numpy as np
 from executorch.backends.nxp.backend.edge_helper import try_get_arg
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    _is_dequant_node,
     CustomDelegationOptions,
     is_not_qdq_node,
     NodeConverter,
 )
 from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import (
     BuiltinOperator,
 )
+from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
+    maximum_options,
+    minimum_options,
+)
 from executorch.backends.nxp.backend.neutron_operator_support import (
     activation_supported_on_target,
 )
@@ -21,6 +30,16 @@
 from torch.nn import Parameter
 
 
+def _is_convertible_to_relu(node):
+    bounds = ClampConverter._get_clamp_bounds(node)
+
+    # Only some specific bounds are supported on the target hardware.
+    if bounds not in ClampConverter.SUPPORTED_BOUNDS.values():
+        return False
+
+    return True
+
+
 class ClampConverter(NodeConverter):
     SUPPORTED_BOUNDS = {
         "ReluN1To1": (-1, 1),
@@ -48,7 +67,7 @@ def _get_clamp_bounds(clamp_node: Node) -> tuple[float | None, float | None]:
     def _is_supported_in_IR(
         node: Node,
         parameters_mapping: dict[str, Parameter],
-        custom_delegation_options: CustomDelegationOptions,
+        _: CustomDelegationOptions,
     ) -> bool:
         # No NeutronIR-specific restrictions.
         return True
@@ -58,22 +77,19 @@ def _is_supported_on_target(
         node: Node,
         neutron_target_spec: NeutronTargetSpec,
         parameters_mapping: dict[str, Parameter],
-        custom_delegation_options: CustomDelegationOptions,
+        _: CustomDelegationOptions,
     ) -> bool:
-        bounds = ClampConverter._get_clamp_bounds(node)
-
-        # Only some specific bounds are supported on the target hardware.
-        if bounds not in ClampConverter.SUPPORTED_BOUNDS.values():
-            return False
+        if neutron_target_spec.use_new_flow_neutron_c:
+            return True
 
-        return True
+        return _is_convertible_to_relu(node)
 
     @classmethod
     def supports_partitioning_result(
         cls,
         node: Node,
         partition_list: list[Partition],
-        custom_delegation_options: CustomDelegationOptions,
+        _: CustomDelegationOptions,
         neutron_target_spec: NeutronTargetSpec,
         parameters_mapping: dict[str, Parameter],
     ) -> bool:
@@ -91,6 +107,15 @@ def supports_partitioning_result(
 
         return True
 
+    @staticmethod
+    def propagate_quantization(from_node, to_node):
+        to_node.quantization = copy(from_node.quantization)
+
+    @staticmethod
+    def _quantize_value(value, zp, scale, quant_min, quant_max):
+        rescaled_value = round(value / scale) + zp
+        return np.clip(rescaled_value, quant_min, quant_max)
+
     def convert(self, node: Node):
         """Convert the `aten.clamp.default` operator to Neutron IR `Relu*` operators.
         The schema is:
@@ -101,13 +126,57 @@ def convert(self, node: Node):
             ) -> Tensor
         """
         self.assert_convertible(node)
+        to_relu = _is_convertible_to_relu(node)
 
         bounds = self._get_clamp_bounds(node)
-
         t_op = self._create_tflite_op_with_io_tensors(node)
 
-        # noinspection PyTypeChecker,PyUnboundLocalVariable
-        t_op.opcode_index = self.builder.op_code_index_for_op_type(
-            self.BOUNDS_TO_NEUTRON_IR_OP[bounds]
+        if not self.neutron_target_spec.use_new_flow_neutron_c or to_relu:
+            # noinspection PyTypeChecker,PyUnboundLocalVariable
+            t_op.opcode_index = self.builder.op_code_index_for_op_type(
+                self.BOUNDS_TO_NEUTRON_IR_OP[bounds]
+            )
+            self.builder.append_operators([t_op])
+            return
+
+        q_node = node.args[0]
+        assert _is_dequant_node(q_node)
+        _, scale, zp, quant_min, quant_max, _ = q_node.args
+
+        x = t_op.tmp_inputs[0]
+        y = t_op.tmp_outputs[0]
+
+        if x.quantization is not None and y.quantization is None:
+            self.propagate_quantization(x, y)
+
+        if x.quantization != y.quantization:
+            raise AssertionError(
+                "Input and output quantization should be same in order to convert to max/min."
+            )
+
+        max_y = self.builder.duplicate_tensor(x)
+
+        min_value, max_value = bounds
+        min_value = self._quantize_value(min_value, zp, scale, quant_min, quant_max)
+        max_value = self._quantize_value(max_value, zp, scale, quant_min, quant_max)
+
+        min_tensor = self.builder.create_tensor_for_data(
+            np.array([min_value], np.int8), "min"
+        )
+        self.propagate_quantization(x, min_tensor)
+        max_tensor = self.builder.create_tensor_for_data(
+            np.array([max_value], np.int8), "max"
         )
-        self.builder.append_operators([t_op])
+        self.propagate_quantization(x, max_tensor)
+
+        max_op = tflite_model.Operator(builtin_options=maximum_options.Maximum())
+        max_op.tmp_inputs = [x, max_tensor]
+        max_op.tmp_outputs = [max_y]
+
+        min_op = tflite_model.Operator(builtin_options=minimum_options.Minimum())
+        min_op.tmp_inputs = [max_y, min_tensor]
+        min_op.tmp_outputs = [y]
+
+        self.propagate_quantization(x, max_y)
+
+        self.builder.append_operators([max_op, min_op])

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

@@ -7,7 +7,6 @@
 
 import numpy as np
 import torch
-
 from executorch.backends.nxp.backend.edge_helper import try_get_arg
 from executorch.backends.nxp.backend.ir.converter.conversion import (
     aten_translator,
@@ -74,7 +73,7 @@ def _is_supported_on_target(
             MaxPool2DWithIndicesConverter._get_node_args(node)
         )
 
-        if custom_delegation_options.use_new_flow_neutron_c:
+        if neutron_target_spec.use_new_flow_neutron_c:
             # Requirements specified by the new Neutron flow documentation.
 
             supported_types = [torch.int8, torch.uint8]

backends/nxp/backend/neutron_target_spec.py

@@ -8,12 +8,10 @@
 from enum import Enum
 
 import torch
-
 from executorch.backends.nxp.backend.neutron_converter_manager import (
     NeutronConverterManager,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
-
 from torch.fx import Node
 
 
@@ -98,13 +96,17 @@ class NeutronTargetSpec:
     The functionality for probing the properties of Neutron Target.
     """
 
-    def __init__(self, target: str):
+    def __init__(self, target: str, use_new_flow_neutron_c: bool = False):
 
         converter_manager = NeutronConverterManager()
         converter_manager.verify_target(target)
         neutron_converter = converter_manager.get_converter()
         self.neutron_target = neutron_converter.getNeutronTarget(target)
 
+        # The new neutron converter flow has different constraints for supported operators. These need to be addressed when
+        # deciding is operator is delegated or not in _is_supported_on_target().
+        self.use_new_flow_neutron_c = use_new_flow_neutron_c
+
         if self.is_subsystem():
             raise ValueError(
                 f"Target `{target}` is not a neutron-C target. Only MCU targets are supported at the moment."

backends/nxp/nxp_backend.py

@@ -14,7 +14,6 @@
 
 import numpy as np
 import torch
-
 from executorch.backends.nxp.backend.custom_delegation_options import (
     CustomDelegationOptions,
 )
@@ -86,7 +85,9 @@ def neutron_compile_spec(
         :return: self for method chaining
         """
 
-        self.config = NeutronTargetSpec(config)
+        self.config = NeutronTargetSpec(
+            config, use_new_flow_neutron_c=use_new_flow_neutron_c
+        )
 
         assert (
             self.output_format is None
@@ -230,11 +231,11 @@ def preprocess(  # noqa C901
             )
             tflite_model, io_formats = EdgeProgramToIRConverter().convert_program(
                 edge_program,
-                neutron_target_spec=NeutronTargetSpec(target),
-                conversion_config=conversion_config,
-                custom_delegation_options=CustomDelegationOptions(
-                    use_new_flow_neutron_c=use_new_flow_neutron_c
+                neutron_target_spec=NeutronTargetSpec(
+                    target, use_new_flow_neutron_c=use_new_flow_neutron_c
                 ),
+                conversion_config=conversion_config,
+                custom_delegation_options=CustomDelegationOptions(),
             )
 
             neutron_model = NeutronConverterManager(dump_kernel_selection_code).convert(