[Backend Tester] Add quantized test flows for XNNPACK and Core ML

ghstack-source-id: 576871d
ghstack-comment-id: 3105090683
Pull-Request: #12733

Author: GregoryComer

backends/apple/coreml/test/tester.py

@@ -4,23 +4,64 @@
 # 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
+from typing import Any, List, Optional, Sequence, Tuple
 
+import coremltools as ct
 import executorch
 import executorch.backends.test.harness.stages as BaseStages
-
+import functools
 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 _get_static_int8_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_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 CoreMLPartitioner(
+                compile_specs=CoreMLBackend.generate_compile_specs(
+                    minimum_deployment_target=minimum_deployment_target
+                )
+            ),
         )
 
 
@@ -29,9 +70,14 @@ 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: CoreMLPartitioner(
+               compile_specs=CoreMLBackend.generate_compile_specs(
+                    minimum_deployment_target=minimum_deployment_target
+                ) 
+            ),
             partitioners=partitioners,
             edge_compile_config=edge_compile_config,
         )
@@ -43,13 +89,15 @@ 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

@@ -31,7 +31,8 @@ def __init__(
         self.calibrate = calibrate
         self.calibration_samples = calibration_samples
 
-        self.quantizer.set_global(self.quantization_config)
+        if self.quantization_config is not None:
+            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)
@@ -21,42 +22,35 @@ 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. """
+    
+    quantize_stage_factory: Callable[..., Quantize] | None = None
+    """ A factory function which instantiates a Quantize stage. Can be None to use the tester's default. """
 
-def create_xnnpack_flow() -> TestFlow | None:
+def all_flows() -> dict[str, TestFlow]:
+    flows = []
+    
     try:
-        from executorch.backends.xnnpack.test.tester import Tester as XnnpackTester
-
-        return TestFlow(
-            name="xnnpack",
-            backend="xnnpack",
-            tester_factory=XnnpackTester,
-        )
-    except Exception:
-        logger.info("Skipping XNNPACK flow registration due to import failure.")
-        return None
-
+        from executorch.backends.test.suite.flows.xnnpack import XNNPACK_TEST_FLOW, XNNPACK_STATIC_INT8_TEST_FLOW
+        flows += [
+            XNNPACK_TEST_FLOW,
+            XNNPACK_STATIC_INT8_TEST_FLOW,
+        ]
+    except Exception as e:
+        logger.info(f"Skipping XNNPACK flow registration: {e}")
 
-def create_coreml_flow() -> TestFlow | None:
     try:
-        from executorch.backends.apple.coreml.test.tester import CoreMLTester
+        from executorch.backends.test.suite.flows.coreml import COREML_TEST_FLOW, COREML_STATIC_INT8_TEST_FLOW
+        flows += [
+            COREML_TEST_FLOW,
+            COREML_STATIC_INT8_TEST_FLOW,
+        ]
+    except Exception as e:
+        logger.info(f"Skipping Core ML flow registration: {e}")
 
-        return TestFlow(
-            name="coreml",
-            backend="coreml",
-            tester_factory=CoreMLTester,
-        )
-    except Exception:
-        logger.info("Skipping Core ML flow registration due to import failure.")
-        return None
-
-
-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,24 @@
+import coremltools
+import functools
+
+from executorch.backends.apple.coreml.test.tester import CoreMLTester
+from executorch.backends.test.suite.flow import TestFlow
+from typing import Any
+
+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,36 @@
+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
+)
+from typing import Callable
+
+import logging
+
+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_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", create_quantize_stage)
+
+XNNPACK_TEST_FLOW = _create_xnnpack_flow()
+XNNPACK_STATIC_INT8_TEST_FLOW = _create_xnnpack_static_int8_flow()

backends/test/suite/models/__init__.py

@@ -49,7 +49,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 +104,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 +117,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,
     )

backends/test/suite/models/test_torchaudio.py

@@ -12,6 +12,7 @@
 import torch
 import torchaudio
 
+from executorch.backends.test.suite.flow import TestFlow
 from executorch.backends.test.suite.models import (
     model_test_cls,
     model_test_params,
@@ -50,7 +51,7 @@ def forward(
 class TorchAudio(unittest.TestCase):
     @model_test_params(dtypes=[torch.float32], supports_dynamic_shapes=False)
     def test_conformer(
-        self, dtype: torch.dtype, use_dynamic_shapes: bool, tester_factory: Callable
+        self, flow: TestFlow, dtype: torch.dtype, use_dynamic_shapes: bool
     ):
         inner_model = torchaudio.models.Conformer(
             input_dim=80,
@@ -70,11 +71,11 @@ def test_conformer(
             encoder_padding_mask,
         )
 
-        run_model_test(model, inputs, dtype, None, tester_factory)
+        run_model_test(model, inputs, flow, dtype, None)
 
     @model_test_params(dtypes=[torch.float32])
     def test_wav2letter(
-        self, dtype: torch.dtype, use_dynamic_shapes: bool, tester_factory: Callable
+        self, flow: TestFlow, dtype: torch.dtype, use_dynamic_shapes: bool
     ):
         model = torchaudio.models.Wav2Letter()
         inputs = (torch.randn(1, 1, 1024, dtype=dtype),)
@@ -87,11 +88,11 @@ def test_wav2letter(
             if use_dynamic_shapes
             else None
         )
-        run_model_test(model, inputs, dtype, dynamic_shapes, tester_factory)
+        run_model_test(model, inputs, flow, dtype, dynamic_shapes)
 
     @unittest.skip("This model times out on all backends.")
     def test_wavernn(
-        self, dtype: torch.dtype, use_dynamic_shapes: bool, tester_factory: Callable
+        self, flow: TestFlow, dtype: torch.dtype, use_dynamic_shapes: bool,
     ):
         model = torchaudio.models.WaveRNN(
             upsample_scales=[5, 5, 8], n_classes=512, hop_length=200
@@ -103,4 +104,4 @@ def test_wavernn(
             torch.randn(1, 1, 128, 64),  # specgram
         )
 
-        run_model_test(model, inputs, dtype, None, tester_factory)
+        run_model_test(model, inputs, flow, dtype, None)