[Backend Tester] Add quantized test flows for XNNPACK and Core ML (#12733)

Add quantized test flows core static int8 quantization for XNNPACK and
Core ML backends. I also ended up doing some light refactoring on the
test signature to pass the TestFlow class into the individual tests.
This is done to allow for passing quantization parameters into the inner
test.

Author: GregoryComer

backends/apple/coreml/test/tester.py

@@ -4,23 +4,73 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, List, Optional, Tuple
+import functools
+from typing import Any, List, Optional, Sequence, Tuple
 
+import coremltools as ct
 import executorch
 import executorch.backends.test.harness.stages as BaseStages
-
 import torch
+
+from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
+from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer
 from executorch.backends.test.harness import Tester as TesterBase
 from executorch.backends.test.harness.stages import StageType
 from executorch.exir import EdgeCompileConfig
 from executorch.exir.backend.partitioner import Partitioner
 
 
+def _create_default_partitioner(
+    minimum_deployment_target: Any = ct.target.iOS15,
+) -> CoreMLPartitioner:
+    return CoreMLPartitioner(
+        compile_specs=CoreMLBackend.generate_compile_specs(
+            minimum_deployment_target=minimum_deployment_target
+        )
+    )
+
+
+def _get_static_int8_linear_qconfig():
+    return ct.optimize.torch.quantization.LinearQuantizerConfig(
+        global_config=ct.optimize.torch.quantization.ModuleLinearQuantizerConfig(
+            quantization_scheme="symmetric",
+            activation_dtype=torch.quint8,
+            weight_dtype=torch.qint8,
+            weight_per_channel=True,
+        )
+    )
+
+
+class Quantize(BaseStages.Quantize):
+    def __init__(
+        self,
+        quantizer: Optional[CoreMLQuantizer] = None,
+        quantization_config: Optional[Any] = None,
+        calibrate: bool = True,
+        calibration_samples: Optional[Sequence[Any]] = None,
+        is_qat: Optional[bool] = False,
+    ):
+        super().__init__(
+            quantizer=quantizer
+            or CoreMLQuantizer(
+                quantization_config or _get_static_int8_linear_qconfig()
+            ),
+            calibrate=calibrate,
+            calibration_samples=calibration_samples,
+            is_qat=is_qat,
+        )
+
+
 class Partition(BaseStages.Partition):
-    def __init__(self, partitioner: Optional[Partitioner] = None):
+    def __init__(
+        self,
+        partitioner: Optional[Partitioner] = None,
+        minimum_deployment_target: Optional[Any] = ct.target.iOS15,
+    ):
         super().__init__(
-            partitioner=partitioner or CoreMLPartitioner,
+            partitioner=partitioner
+            or _create_default_partitioner(minimum_deployment_target),
         )
 
 
@@ -29,9 +79,12 @@ def __init__(
         self,
         partitioners: Optional[List[Partitioner]] = None,
         edge_compile_config: Optional[EdgeCompileConfig] = None,
+        minimum_deployment_target: Optional[Any] = ct.target.iOS15,
     ):
         super().__init__(
-            default_partitioner_cls=CoreMLPartitioner,
+            default_partitioner_cls=lambda: _create_default_partitioner(
+                minimum_deployment_target
+            ),
             partitioners=partitioners,
             edge_compile_config=edge_compile_config,
         )
@@ -43,13 +96,20 @@ def __init__(
         module: torch.nn.Module,
         example_inputs: Tuple[torch.Tensor],
         dynamic_shapes: Optional[Tuple[Any]] = None,
+        minimum_deployment_target: Optional[Any] = ct.target.iOS15,
     ):
         # Specialize for XNNPACK
         stage_classes = (
             executorch.backends.test.harness.Tester.default_stage_classes()
             | {
-                StageType.PARTITION: Partition,
-                StageType.TO_EDGE_TRANSFORM_AND_LOWER: ToEdgeTransformAndLower,
+                StageType.QUANTIZE: Quantize,
+                StageType.PARTITION: functools.partial(
+                    Partition, minimum_deployment_target=minimum_deployment_target
+                ),
+                StageType.TO_EDGE_TRANSFORM_AND_LOWER: functools.partial(
+                    ToEdgeTransformAndLower,
+                    minimum_deployment_target=minimum_deployment_target,
+                ),
             }
         )
 

backends/test/harness/stages/quantize.py

@@ -25,13 +25,15 @@ def __init__(
         calibrate: bool = True,
         calibration_samples: Optional[Sequence[Any]] = None,
         is_qat: Optional[bool] = False,
+        set_global: bool = True,
     ):
         self.quantizer = quantizer
         self.quantization_config = quantization_config
         self.calibrate = calibrate
         self.calibration_samples = calibration_samples
 
-        self.quantizer.set_global(self.quantization_config)
+        if self.quantization_config is not None and set_global:
+            self.quantizer.set_global(self.quantization_config)
 
         self.converted_graph = None
         self.is_qat = is_qat

backends/test/harness/tester.py

@@ -1,6 +1,6 @@
 import random
 from collections import Counter, OrderedDict
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
 
@@ -33,7 +33,7 @@ def __init__(
         self,
         module: torch.nn.Module,
         example_inputs: Tuple[torch.Tensor],
-        stage_classes: Dict[StageType, Type],
+        stage_classes: Dict[StageType, Callable],
         dynamic_shapes: Optional[Tuple[Any]] = None,
     ):
         module.eval()
@@ -81,7 +81,7 @@ def __init__(
         self.stage_output = None
 
     @staticmethod
-    def default_stage_classes() -> Dict[StageType, Type]:
+    def default_stage_classes() -> Dict[StageType, Callable]:
         """
         Returns a map of StageType to default Stage implementation.
         """

backends/test/suite/__init__.py

@@ -129,7 +129,7 @@ def _make_wrapped_test(
     def wrapped_test(self):
         with TestContext(test_name, flow.name, params):
             test_kwargs = params or {}
-            test_kwargs["tester_factory"] = flow.tester_factory
+            test_kwargs["flow"] = flow
 
             test_func(self, **test_kwargs)
 
@@ -175,7 +175,7 @@ def load_tests(loader, suite, pattern):
 
 
 class OperatorTest(unittest.TestCase):
-    def _test_op(self, model, inputs, tester_factory):
+    def _test_op(self, model, inputs, flow: TestFlow):
         context = get_active_test_context()
 
         # This should be set in the wrapped test. See _make_wrapped_test above.
@@ -184,9 +184,8 @@ def _test_op(self, model, inputs, tester_factory):
         run_summary = run_test(
             model,
             inputs,
-            tester_factory,
+            flow,
             context.test_name,
-            context.flow_name,
             context.params,
         )
 

backends/test/suite/flow.py

@@ -1,9 +1,10 @@
 import logging
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Callable
 
 from executorch.backends.test.harness import Tester
+from executorch.backends.test.harness.stages import Quantize
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -22,41 +23,43 @@ class TestFlow:
     backend: str
     """ The name of the target backend. """
 
-    tester_factory: Callable[[], Tester]
+    tester_factory: Callable[..., Tester]
     """ A factory function that returns a Tester instance for this lowering flow. """
 
+    quantize: bool = field(default=False)
+    """ Whether to tester should run the quantize stage on the model. """
 
-def create_xnnpack_flow() -> TestFlow | None:
-    try:
-        from executorch.backends.xnnpack.test.tester import Tester as XnnpackTester
+    quantize_stage_factory: Callable[..., Quantize] | None = None
+    """ A factory function which instantiates a Quantize stage. Can be None to use the tester's default. """
 
-        return TestFlow(
-            name="xnnpack",
-            backend="xnnpack",
-            tester_factory=XnnpackTester,
-        )
-    except Exception:
-        logger.info("Skipping XNNPACK flow registration due to import failure.")
-        return None
 
+def all_flows() -> dict[str, TestFlow]:
+    flows = []
 
-def create_coreml_flow() -> TestFlow | None:
     try:
-        from executorch.backends.apple.coreml.test.tester import CoreMLTester
+        from executorch.backends.test.suite.flows.xnnpack import (
+            XNNPACK_STATIC_INT8_PER_CHANNEL_TEST_FLOW,
+            XNNPACK_TEST_FLOW,
+        )
 
-        return TestFlow(
-            name="coreml",
-            backend="coreml",
-            tester_factory=CoreMLTester,
+        flows += [
+            XNNPACK_TEST_FLOW,
+            XNNPACK_STATIC_INT8_PER_CHANNEL_TEST_FLOW,
+        ]
+    except Exception as e:
+        logger.info(f"Skipping XNNPACK flow registration: {e}")
+
+    try:
+        from executorch.backends.test.suite.flows.coreml import (
+            COREML_STATIC_INT8_TEST_FLOW,
+            COREML_TEST_FLOW,
         )
-    except Exception:
-        logger.info("Skipping Core ML flow registration due to import failure.")
-        return None
 
+        flows += [
+            COREML_TEST_FLOW,
+            COREML_STATIC_INT8_TEST_FLOW,
+        ]
+    except Exception as e:
+        logger.info(f"Skipping Core ML flow registration: {e}")
 
-def all_flows() -> dict[str, TestFlow]:
-    flows = [
-        create_xnnpack_flow(),
-        create_coreml_flow(),
-    ]
     return {f.name: f for f in flows if f is not None}

backends/test/suite/flows/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# 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.
+
+# pyre-unsafe

backends/test/suite/flows/coreml.py

@@ -0,0 +1,30 @@
+import functools
+from typing import Any
+
+import coremltools
+
+from executorch.backends.apple.coreml.test.tester import CoreMLTester
+from executorch.backends.test.suite.flow import TestFlow
+
+
+def _create_coreml_flow(
+    name: str,
+    quantize: bool = False,
+    minimum_deployment_target: Any = coremltools.target.iOS15,
+) -> TestFlow:
+    return TestFlow(
+        name,
+        backend="coreml",
+        tester_factory=functools.partial(
+            CoreMLTester, minimum_deployment_target=minimum_deployment_target
+        ),
+        quantize=quantize,
+    )
+
+
+COREML_TEST_FLOW = _create_coreml_flow("coreml")
+COREML_STATIC_INT8_TEST_FLOW = _create_coreml_flow(
+    "coreml_static_int8",
+    quantize=True,
+    minimum_deployment_target=coremltools.target.iOS17,
+)

backends/test/suite/flows/xnnpack.py

@@ -0,0 +1,49 @@
+import logging
+from typing import Callable
+
+from executorch.backends.test.harness.stages import Quantize
+from executorch.backends.test.suite.flow import TestFlow
+from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
+    get_symmetric_quantization_config,
+)
+from executorch.backends.xnnpack.test.tester import (
+    Quantize as XnnpackQuantize,
+    Tester as XnnpackTester,
+)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+def _create_xnnpack_flow_base(
+    name: str, quantize_stage_factory: Callable[..., Quantize] | None = None
+) -> TestFlow:
+    return TestFlow(
+        name,
+        backend="xnnpack",
+        tester_factory=XnnpackTester,
+        quantize=quantize_stage_factory is not None,
+        quantize_stage_factory=quantize_stage_factory,
+    )
+
+
+def _create_xnnpack_flow() -> TestFlow:
+    return _create_xnnpack_flow_base("xnnpack")
+
+
+def _create_xnnpack_static_int8_per_channel_flow() -> TestFlow:
+    def create_quantize_stage() -> Quantize:
+        qparams = get_symmetric_quantization_config(is_per_channel=True)
+        return XnnpackQuantize(
+            quantization_config=qparams,
+        )
+
+    return _create_xnnpack_flow_base(
+        "xnnpack_static_int8_per_channel", create_quantize_stage
+    )
+
+
+XNNPACK_TEST_FLOW = _create_xnnpack_flow()
+XNNPACK_STATIC_INT8_PER_CHANNEL_TEST_FLOW = (
+    _create_xnnpack_static_int8_per_channel_flow()
+)

backends/test/suite/models/__init__.py

@@ -12,7 +12,6 @@
 from typing import Any, Callable
 
 import torch
-from executorch.backends.test.harness import Tester
 from executorch.backends.test.suite import get_test_flows
 from executorch.backends.test.suite.context import get_active_test_context, TestContext
 from executorch.backends.test.suite.flow import TestFlow
@@ -49,7 +48,7 @@ def wrapped_test(self):
             "use_dynamic_shapes": use_dynamic_shapes,
         }
         with TestContext(test_name, flow.name, params):
-            test_func(self, dtype, use_dynamic_shapes, flow.tester_factory)
+            test_func(self, flow, dtype, use_dynamic_shapes)
 
     dtype_name = str(dtype)[6:]  # strip "torch."
     test_name = f"{test_func.__name__}_{flow.name}_{dtype_name}"
@@ -104,9 +103,9 @@ def inner_decorator(func: Callable) -> Callable:
 def run_model_test(
     model: torch.nn.Module,
     inputs: tuple[Any],
+    flow: TestFlow,
     dtype: torch.dtype,
     dynamic_shapes: Any | None,
-    tester_factory: Callable[[], Tester],
 ):
     model = model.to(dtype)
     context = get_active_test_context()
@@ -117,9 +116,8 @@ def run_model_test(
     run_summary = run_test(
         model,
         inputs,
-        tester_factory,
+        flow,
         context.test_name,
-        context.flow_name,
         context.params,
         dynamic_shapes=dynamic_shapes,
     )