Arm backend: Add label attribute to QuantizationConfig

The quantizer reporter logs the quantization config in a human-readable
format. Prior to this patch, this was done with the help of a dict
called `SUPPORTED_QCONFIGS`, which was defined in quantizer_reporter.py
and populated by the user. This patch reworks this concept by instead
adding a label attribute to `QuantizationConfig` that the reporter can
use to print the config in a human-readable format.

Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com>
Change-Id: I38e80c9c3d57fb9d858119fe4281b713bf472475

Author: Martin Lindström

backends/arm/quantizer/arm_quantizer.py

@@ -41,7 +41,6 @@
 
 from executorch.backends.cortex_m.quantizer_reporter import (
     QuantizerReporter,
-    SUPPORTED_QCONFIGS,
     SUPPORTED_QSPECS,
 )
 
@@ -219,20 +218,28 @@ def get_symmetric_quantization_config(
         bias_quantization_spec = _get_int32_bias_qspec
 
     if is_dynamic:
-        quantization_config = TOSAQuantizationConfig(
-            act_quantization_spec,
-            None,
-            weight_quantization_spec,
-            bias_quantization_spec,
-        )
+        output_activation = None
     else:
-        quantization_config = TOSAQuantizationConfig(
-            act_quantization_spec,
-            act_quantization_spec,
-            weight_quantization_spec,
-            bias_quantization_spec,
-        )
-    return quantization_config
+        output_activation = act_quantization_spec
+
+    module_name = __name__.rsplit(".", maxsplit=1)[-1]
+    label = (
+        f"{module_name}.get_symmetric_quantization_config("
+        f"per_channel={int(is_per_channel)}, "
+        f"qat={int(is_qat)}, "
+        f"dynamic={int(is_dynamic)}, "
+        f"act_range=[{act_qmin}, {act_qmax}], "
+        f"weight_range=[{weight_qmin}, {weight_qmax}]"
+        ")"
+    )
+
+    return TOSAQuantizationConfig(
+        act_quantization_spec,
+        output_activation,
+        weight_quantization_spec,
+        bias_quantization_spec,
+        label,
+    )
 
 
 @functools.lru_cache
@@ -357,22 +364,31 @@ def get_symmetric_a16w8_quantization_config(
         is_qat=is_qat,
         is_dynamic=is_dynamic,
     )
-    # Replace activation quantization spec with 16-bit version
+
     if is_dynamic:
-        quantization_config = TOSAQuantizationConfig(
-            act_quantization_spec,  # 16-bit input activations
-            None,
-            base_config.weight,  # 8-bit weights from base config
-            base_config.bias,  # bias from base config
-        )
+        output_activation = None
     else:
-        quantization_config = TOSAQuantizationConfig(
-            act_quantization_spec,  # 16-bit input activations
-            act_quantization_spec,  # 16-bit output activations
-            base_config.weight,  # 8-bit weights from base config
-            base_config.bias,  # bias from base config
-        )
-    return quantization_config
+        output_activation = act_quantization_spec
+
+    module_name = __name__.rsplit(".", maxsplit=1)[-1]
+    label = (
+        f"{module_name}.get_symmetric_a16w8_quantization_config("
+        f"per_channel={int(is_per_channel)}, "
+        f"qat={int(is_qat)}, "
+        f"dynamic={int(is_dynamic)}, "
+        f"act_range=[{act_quantization_spec.quant_min}, {act_quantization_spec.quant_max}], "
+        f"weight_range=[{weight_qmin}, {weight_qmax}]"
+        ")"
+    )
+
+    # Replace activation quantization spec with 16-bit version
+    return TOSAQuantizationConfig(
+        act_quantization_spec,  # 16-bit input activations
+        output_activation,
+        base_config.weight,  # 8-bit weights from base config
+        base_config.bias,  # bias from base config
+        label,
+    )
 
 
 # Register supported quantization configs and qspecs in the reporter for human-readable reporting
@@ -389,16 +405,6 @@ def get_symmetric_a16w8_quantization_config(
 _symmetric_a16w8_config_per_tensor = get_symmetric_a16w8_quantization_config(
     is_per_channel=False
 )
-SUPPORTED_QCONFIGS.update(
-    {
-        _symmetric_a8w8_config_per_channel: f"{__name__}.get_symmetric_quantization_config(is_per_channel=True)",
-        _symmetric_a16w8_config_per_channel: f"{__name__}.get_symmetric_a16w8_quantization_config(is_per_channel=True)",
-        _symmetric_a8w4_config_per_channel: f"{__name__}.get_symmetric_a8w4_quantization_config(is_per_channel=True)",
-        _symmetric_a8w8_config_per_tensor: f"{__name__}.get_symmetric_quantization_config(is_per_channel=False)",
-        _symmetric_a16w8_config_per_tensor: f"{__name__}.get_symmetric_a16w8_quantization_config(is_per_channel=False)",
-        _symmetric_a8w4_config_per_tensor: f"{__name__}.get_symmetric_a8w4_quantization_config(is_per_channel=False)",
-    }
-)
 
 SUPPORTED_QSPECS.update(
     {

backends/arm/quantizer/arm_quantizer_utils.py

@@ -21,7 +21,10 @@
 from executorch.backends.arm.common.annotation_meta import ArmAnnotationInfo
 from executorch.backends.arm.constants import DISALLOW_TFA_META_KEY
 from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
-from executorch.backends.cortex_m.quantizer_reporter import QuantizerReporterUser
+from executorch.backends.cortex_m.quantizer_reporter import (
+    QuantizerInfo,
+    QuantizerReporterUser,
+)
 from torch.fx import Node
 
 from torchao.quantization.pt2e.quantizer import (
@@ -253,22 +256,19 @@ def __init__(
         self.pattern_matcher: "PatternMatcher" = pattern_matcher
 
     def get_quantizer_info(self):
-        from executorch.backends.cortex_m.quantizer_reporter import (
-            QuantizerInfo,
-            SUPPORTED_QCONFIGS,
-        )
-
         name = self.__class__.__name__
         targeted_nodes_description = str(self.node_finder)
-        quantization_config_path = SUPPORTED_QCONFIGS.get(
-            self.quantization_config, "UNREGISTERED_QCONFIG"
+        qconfig_label = (
+            self.quantization_config.label
+            if self.quantization_config.label is not None
+            else self.quantization_config.__class__.__name__  # no label, fallback to class name
         )
         support_config_path = self.pattern_matcher.support_dict_name
 
         return QuantizerInfo(
             name,
             targeted_nodes_description,
-            quantization_config_path,
+            qconfig_label,
             support_config_path,
         )
 
@@ -490,16 +490,14 @@ def __init__(self, targets: Optional[list[Callable[..., object]]] = None) -> Non
             )
 
     def get_quantizer_info(self):
-        from executorch.backends.cortex_m.quantizer_reporter import QuantizerInfo
-
         name = self.__class__.__name__
         targeted_nodes_description = ""
-        quantization_config_path = "SHARED_QCONFIG"
+        qconfig_label = "shared qparams for connected targeted nodes"
         support_config_path = self.support_config_path
         return QuantizerInfo(
             name,
             targeted_nodes_description,
-            quantization_config_path,
+            qconfig_label,
             support_config_path,
         )
 

backends/arm/quantizer/quantization_config.py

@@ -46,6 +46,7 @@ class QuantizationConfig:
     output_activation: Optional[QuantizationSpecBase]
     weight: Optional[QuantizationSpecBase]
     bias: Optional[QuantizationSpecBase] | Callable[[Any], Any]
+    label: Optional[str] = None  # Optional label for debugging/visualization purposes
 
     def get_input_act_qspec(
         self, node: Optional[Node] = None, input_node: Optional[Node] = None

backends/cortex_m/quantizer/quantization_configs.py

@@ -10,10 +10,7 @@
     _get_int32_per_channel_bias_qspec,
 )
 from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
-from executorch.backends.cortex_m.quantizer_reporter import (
-    SUPPORTED_QCONFIGS,
-    SUPPORTED_QSPECS,
-)
+from executorch.backends.cortex_m.quantizer_reporter import SUPPORTED_QSPECS
 from torch.fx import Node
 from torchao.quantization.pt2e import (
     HistogramObserver,
@@ -156,6 +153,7 @@ def get_bias_qspec(
     INT8_ACTIVATION_PER_TENSOR_QSPEC,
     INT8_WEIGHT_PER_TENSOR_QSPEC,
     _get_int32_bias_qspec,
+    f"{__name__}.INT8_PER_TENSOR_CONFIG",
 )
 
 
@@ -164,18 +162,10 @@ def get_bias_qspec(
     INT8_ACTIVATION_PER_TENSOR_QSPEC,
     INT8_WEIGHT_PER_CHANNEL_QSPEC,
     _get_int32_per_channel_bias_qspec,
+    f"{__name__}.INT8_PER_CHANNEL_CONFIG",
 )
 
 
-# Register supported quantization configs and qspecs in the reporter for human-readable reporting
-# MLETORCH-1854: Temporary solution, refactor to automatically register these instead
-SUPPORTED_QCONFIGS.update(
-    {
-        INT8_PER_CHANNEL_CONFIG: f"{__name__}.INT8_PER_CHANNEL_QCONFIG",
-        INT8_PER_TENSOR_CONFIG: f"{__name__}.INT8_PER_TENSOR_QCONFIG",
-    }
-)
-
 SUPPORTED_QSPECS.update(
     {
         INT8_ACTIVATION_PER_TENSOR_QSPEC: "INT8_ACTIVATION_PER_TENSOR_QSPEC",

backends/cortex_m/quantizer_reporter.py

@@ -17,7 +17,7 @@
 
 import logging
 from importlib import import_module
-from typing import Any, Callable, cast, Dict, List, NamedTuple, Optional
+from typing import Callable, cast, Dict, List, NamedTuple, Optional
 
 from torch.fx import GraphModule, Node
 from torchao.quantization.pt2e.quantizer import (
@@ -32,8 +32,7 @@
 logger = logging.getLogger(__name__)
 tabulate = cast(Callable[..., str], import_module("tabulate").tabulate)
 
-# Look-up dicts used to get human readable names for supported quantization configs and specs
-SUPPORTED_QCONFIGS: dict[Any, str] = {}
+# Look-up dicts used to get human readable names for supported quantization specs
 SUPPORTED_QSPECS: dict[QuantizationSpecBase | None, str] = {}
 
 
@@ -77,7 +76,7 @@ class QuantizerInfo(NamedTuple):
 
     name: str
     targeted_nodes_description: str
-    quantization_config_path: str
+    qconfig_label: str
     support_config_path: str
 
 
@@ -112,8 +111,8 @@ class QuantizerReport:
 
     _PREVIOUS_ANNOTATION_REJECT_REASON = "Tried annotating already quantized node."
 
-    def __init__(self, quantizer):
-        self.quantizer = quantizer.get_quantizer_info()
+    def __init__(self, quantizer_info: QuantizerInfo):
+        self.quantizer_info = quantizer_info
         self.accepted_patterns: List[AnnotatedPatternReport] = []
         self.rejected_patterns: List[RejectedPatternReport] = []
 
@@ -180,11 +179,11 @@ def report_reject(self, pattern, reason):
     def get_quantizer_info_rows(self) -> List[str]:
         rows = []
         rows.append(
-            f"{self.quantizer.name} using {self.quantizer.targeted_nodes_description}"
+            f"{self.quantizer_info.name} using {self.quantizer_info.targeted_nodes_description}"
         )
-        rows.append(f"Annotating with {self.quantizer.quantization_config_path}")
+        rows.append(f"Annotating with {self.quantizer_info.qconfig_label}")
         rows.append(
-            f"Supported operators and patterns defined by {self.quantizer.support_config_path}"
+            f"Supported operators and patterns defined by {self.quantizer_info.support_config_path}"
         )
 
         if (
@@ -317,7 +316,7 @@ def set_quantizers(self, quantizers: List[QuantizerReporterUser]) -> None:
                     f"Quantizer {quantizer.__class__.__name__} does not implement QuantizerReporterUser interface and will not report quantization decisions."
                 )
 
-            self.quantizers[quantizer] = QuantizerReport(quantizer)
+            self.quantizers[quantizer] = QuantizerReport(quantizer.get_quantizer_info())
 
     def report_reject(
         self, quantizer: QuantizerReporterUser, pattern: List[Node], reason: str

backends/cortex_m/test/misc/test_quantizer_reporter.py

@@ -38,7 +38,7 @@ def get_quantizer_info(self) -> QuantizerInfo:
         return QuantizerInfo(
             name="DummyQuantizer",
             targeted_nodes_description="dummy nodes",
-            quantization_config_path="dummy.config",
+            qconfig_label="dummy.config",
             support_config_path="dummy.support",
         )