Quantize Addmm, Conv2d, Linear, Mm together with fusable activations

+ Move fused activations to separate QDQ cluster

Author: roman-janik-nxp

backends/nxp/backend/edge_helper.py

@@ -4,10 +4,32 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
+
+from executorch.exir.dialects._ops import ops as exir_ops
+
 from torch.fx import GraphModule, Node
 from torch.nn import Parameter
 
 
+QUANTIZE_OPERATORS = [
+    exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
+    exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+]
+
+DEQUANTIZE_OPERATORS = [
+    exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
+    exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+]
+
+
+def _is_dequantize(node_: Node) -> bool:
+    return node_.op == "call_function" and node_.target in DEQUANTIZE_OPERATORS
+
+
+def _is_quantize(node_: Node) -> bool:
+    return node_.op == "call_function" and node_.target in QUANTIZE_OPERATORS
+
+
 def input_tensor(node: Node, input_index: int) -> torch.Tensor:
     if len(node.all_input_nodes) <= input_index:
         raise IndexError
@@ -62,12 +84,6 @@ def node_is_effectively_static_tensor(
     if node_is_static_tensor(node, parameters_mapping):
         return True
 
-    def _is_dequantize(node_: Node) -> bool:
-        return node_.target.__name__ in {
-            "quantized_decomposed.dequantize_per_tensor.default",
-            "quantized_decomposed.dequantize_per_channel.default",
-        }
-
     return _is_dequantize(node) and node_is_static_tensor(
         node.args[0], parameters_mapping
     )

backends/nxp/backend/neutron_target_spec.py

@@ -7,16 +7,92 @@
 
 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
 
 
 class NeutronHWVersion(Enum):
     N1 = 1
     N3 = 2
 
 
+class NeutronTargetNeutronC:
+    @staticmethod
+    def is_supported_fused_activation__aten(node_: Node) -> bool:
+        """Node operator is supported fused activation on Neutron for Linear and Conv2D."""
+        return node_.op == "call_function" and (
+            node_.target
+            in (
+                torch.ops.aten.relu.default,  # TODO Add torch.ops.aten.leaky_relu.default once it is supported
+                torch.ops.aten.relu_.default,
+                torch.ops.aten.sigmoid.default,
+                torch.ops.aten.sigmoid_.default,
+                torch.ops.aten.tanh.default,
+                torch.ops.aten.tanh_.default,
+            )
+            or (
+                (
+                    node_.target == torch.ops.aten.hardtanh.default
+                    or node_.target == torch.ops.aten.hardtanh_.default
+                )
+                and (
+                    node_.args[1:3] == (0.0, 6.0)  # is converted to Relu6
+                    or node_.args[1:3] == (0.0, float("inf"))  # is converted to Relu
+                )
+            )
+        )
+
+    @staticmethod
+    def is_supported_fused_activation__edge(node_: Node) -> bool:
+        """Node operator is supported fused activation on Neutron for Linear and Conv2D."""
+        return node_.op == "call_function" and (
+            node_.target
+            in (
+                exir_ops.edge.aten.relu.default,  # TODO Add torch.ops.aten.leaky_relu.default once it is supported
+                exir_ops.edge.aten.sigmoid.default,
+                exir_ops.edge.aten.tanh.default,
+            )
+            or (
+                (node_.target == exir_ops.edge.aten.hardtanh.default)
+                and (
+                    node_.args[1:3] == (0.0, 6.0)  # is converted to Relu6
+                    or node_.args[1:3] == (0.0, float("inf"))  # is converted to Relu
+                )
+            )
+        )
+
+    @staticmethod
+    def is_fusable_conv_or_linear__aten(node_: Node) -> bool:
+        """Node operator is supported fusable Linear or Conv2D on Neutron."""
+        return node_.op == "call_function" and (
+            node_.target == torch.ops.aten.conv2d.default
+            or node_.target == torch.ops.aten.addmm.default
+            or node_.target == torch.ops.aten.mm.default
+            or (
+                node_.target == torch.ops.aten.linear.default
+                and len(node_.meta["val"].shape) == 2
+            )
+        )
+
+    @staticmethod
+    def is_fusable_conv_or_linear__edge(node_: Node) -> bool:
+        """Node operator in edge dialect is supported fusable Linear or Conv2D on Neutron."""
+        return node_.op == "call_function" and (
+            node_.target == exir_ops.edge.aten.addmm.default
+            or node_.target == exir_ops.edge.aten.mm.default
+            or (
+                node_.target == exir_ops.edge.aten.convolution.default
+                and len(node_.meta["val"].shape) == 4
+            )
+        )
+
+
 class NeutronTargetSpec:
     """
     The functionality for probing the properties of Neutron Target.
@@ -39,6 +115,9 @@ def __init__(self, target: str, neutron_converter_flavor: str):
                 f"Target `{target}` contains unsupported HW version. Only N3/N3+ targets are supported at the moment."
             )
 
+        # Now only Neutron-C is supported
+        self.neutron_target_info = NeutronTargetNeutronC()
+
     # Target name.
     def get_name(self) -> str:
         return self.neutron_target.name

backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py

@@ -15,6 +15,11 @@
 AddMM = exir_ops.edge.aten.addmm.default
 ViewCopy = exir_ops.edge.aten.view_copy.default
 MM = exir_ops.edge.aten.mm.default
+Conv = exir_ops.edge.aten.convolution.default
+HardTanh = exir_ops.edge.aten.hardtanh.default
+Relu = exir_ops.edge.aten.relu.default
+Sigmoid = exir_ops.edge.aten.sigmoid.default
+Tanh = exir_ops.edge.aten.tanh.default
 
 
 def insert_qdq_pair_after_node(
@@ -175,9 +180,23 @@ class MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass):
     main_cluster_node_to_auxiliary_nodes = {
         AddMM: [
             ViewCopy,
+            HardTanh,
+            Relu,
+            Sigmoid,
+            Tanh,
         ],
         MM: [
             ViewCopy,
+            HardTanh,
+            Relu,
+            Sigmoid,
+            Tanh,
+        ],
+        Conv: [
+            HardTanh,
+            Relu,
+            Sigmoid,
+            Tanh,
         ],
     }
 

backends/nxp/neutron_partitioner.py

@@ -15,6 +15,10 @@
 from executorch.backends.nxp.backend.custom_delegation_options import (
     CustomDelegationOptions,
 )
+from executorch.backends.nxp.backend.edge_helper import (
+    DEQUANTIZE_OPERATORS,
+    QUANTIZE_OPERATORS,
+)
 from executorch.backends.nxp.backend.edge_program_converter import (
     EdgeProgramToIRConverter,
 )
@@ -66,32 +70,26 @@ class QDQCluster:
         compute_node: torch.fx.Node
         ops: list[torch.fx.Node]
 
-    QUANTIZE_OPERATORS = [
-        exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
-        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
-    ]
-
-    DEQUANTIZE_OPERATORS = [
-        exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
-        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
-    ]
-
     AUXILIARY_OPS = [
         operator.getitem,
         exir_ops.edge.aten.view_copy.default,
         exir_ops.edge.aten.permute_copy.default,
+        exir_ops.edge.aten.hardtanh.default,
+        exir_ops.edge.aten.relu.default,
+        exir_ops.edge.aten.sigmoid.default,
+        exir_ops.edge.aten.tanh.default,
     ]
 
     def __init__(self):
         self.cluster_map: dict[str, QDQClusterRecognizer.QDQCluster] = {}
 
     @staticmethod
     def is_quant_node(node: torch.fx.Node) -> bool:
-        return node.target in QDQClusterRecognizer.QUANTIZE_OPERATORS
+        return node.target in QUANTIZE_OPERATORS
 
     @staticmethod
     def is_dequant_node(node: torch.fx.Node) -> bool:
-        return node.target in QDQClusterRecognizer.DEQUANTIZE_OPERATORS
+        return node.target in DEQUANTIZE_OPERATORS
 
     @staticmethod
     def is_auxiliary_node(node: torch.fx.Node) -> bool:
@@ -308,18 +306,17 @@ class NeutronPartitioner(Partitioner):
     def __init__(
         self,
         compile_spec: list[CompileSpec],
+        neutron_target_spec: NeutronTargetSpec,
         custom_delegation_options: CustomDelegationOptions | None = None,
     ) -> None:
         self.delegation_spec = DelegationSpec(NeutronBackend.__name__, compile_spec)
         self.custom_delegation_options = (
             custom_delegation_options or CustomDelegationOptions()
         )
-        target = self.delegation_spec[1][2].value.decode()
-        converter_flavor = self.delegation_spec[1][3].value.decode()
-        self.neutron_target_spec = NeutronTargetSpec(target, converter_flavor)
+        self.neutron_target_spec = neutron_target_spec
 
+    @staticmethod
     def validate_partitioning_result(
-        self,
         graph: Graph,
         partition_list: list[Partition],
         custom_delegation_options: CustomDelegationOptions,

backends/nxp/quantizer/neutron_quantizer.py

@@ -5,10 +5,11 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-
 from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import (
     NeutronAtenPassManager,
 )
+
+from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 from executorch.backends.nxp.quantizer.patterns import (
     AbsPattern,
     AdaptiveAvgPoolPattern,
@@ -181,27 +182,28 @@ def get_supported_operators(cls) -> list[OperatorConfig]:
 
 
 class NeutronQuantizer(ComposableQuantizer):
-    def __init__(self):
+    def __init__(self, neutron_target_spec: NeutronTargetSpec):
+        self.neutron_target_spec = neutron_target_spec
         static_qconfig = QuantizationConfig(act_qspec, act_qspec, wgt_qspec, None)
         static_fc_qconfig = QuantizationConfig(act_qspec, act_qspec, wgt_fc_qspec, None)
         super().__init__(
             [
                 NeutronAtenQuantizer(AbsPattern(), static_qconfig),
                 NeutronAtenQuantizer(AdaptiveAvgPoolPattern(), static_qconfig),
                 NeutronAtenQuantizer(AddTensorPattern(), static_qconfig),
-                NeutronAtenQuantizer(AddmmPattern(), static_fc_qconfig),
+                NeutronAtenQuantizer(AddmmPattern(self), static_fc_qconfig),
                 NeutronAtenQuantizer(AvgPoolPattern(), static_qconfig),
                 NeutronAtenQuantizer(CatPattern(), static_qconfig),
                 NeutronAtenQuantizer(Conv1dPattern(), static_qconfig),
-                NeutronAtenQuantizer(Conv2dPattern(), static_qconfig),
+                NeutronAtenQuantizer(Conv2dPattern(self), static_qconfig),
                 NeutronAtenQuantizer(DropoutPattern(), static_qconfig),
                 NeutronAtenQuantizer(FlattenPattern(), static_qconfig),
                 NeutronAtenQuantizer(HardTanhPattern(), static_qconfig),
                 NeutronAtenQuantizer(HardTanhInPlacePattern(), static_qconfig),
-                NeutronAtenQuantizer(LinearPattern(), static_fc_qconfig),
+                NeutronAtenQuantizer(LinearPattern(self), static_fc_qconfig),
                 NeutronAtenQuantizer(MaxPoolPattern(), static_qconfig),
                 NeutronAtenQuantizer(MeanDimPattern(), static_qconfig),
-                NeutronAtenQuantizer(MmPattern(), static_qconfig),
+                NeutronAtenQuantizer(MmPattern(self), static_qconfig),
                 NeutronAtenQuantizer(PadPattern(), static_qconfig),
                 NeutronAtenQuantizer(PermutePattern(), static_qconfig),
                 NeutronAtenQuantizer(ReluPattern(), static_qconfig),