diff --git a/backends/apple/coreai/quantizer/__init__.py b/backends/apple/coreai/quantizer/__init__.py new file mode 100644 index 00000000000..3c216b0b994 --- /dev/null +++ b/backends/apple/coreai/quantizer/__init__.py @@ -0,0 +1,11 @@ +# 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. + +from executorch.backends.apple.coreai.quantizer.quantizer import CoreAIQuantizer + +__all__ = [ + "CoreAIQuantizer", +] diff --git a/backends/apple/coreai/quantizer/quantizer.py b/backends/apple/coreai/quantizer/quantizer.py new file mode 100644 index 00000000000..0843ddd241a --- /dev/null +++ b/backends/apple/coreai/quantizer/quantizer.py @@ -0,0 +1,87 @@ +# 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. + +"""Core AI quantization front-end for ExecuTorch. + +Keeps the canonical ExecuTorch quantization pipeline:: + + prepare -> calibrate -> convert -> export -> to_edge_transform_and_lower + +``CoreAIQuantizer.convert()`` produces the export-ready, fully quantized graph +(Core AI ``coreai::`` ops). It intentionally runs the entire ``coreai_opt`` +``finalize(CoreAI)`` (``convert_pt2e`` + conv/bn fold + Q/DQ to ``coreai::`` +rewrite + kv-cache relocation), because that whole rewrite must run before +``torch.export``: a graph still carrying ``coreai_opt`` fake-quant cannot be +strict-``torch.export``ed (its ``FakeQuantize.forward`` has a data-dependent +guard). Nothing quant-related is left for the backend ``preprocess`` to do; it +simply lowers the ``coreai::`` ops via ``add_exported_program``. + +``coreai_opt`` (the ``coreai-optimization`` package) is imported lazily so this +module stays importable in environments without it installed. +""" + +from __future__ import annotations + +from typing import Any, Optional, Sequence + +from torch import fx, nn + + +class CoreAIQuantizer: + """PT2E quantization front-end for Core AI, wrapping ``coreai_opt``. + + Standard PT2E lifecycle: :meth:`prepare` (exports + inserts fake-quant), + :meth:`calibration_mode` / :meth:`training_mode`, then :meth:`convert`. + + Unlike stock PT2E (where ``convert`` yields a backend-agnostic Q/DQ graph), + :meth:`convert` here returns the Core AI quantized graph (``coreai::`` ops), + because ``coreai_opt``'s fake-quant has no meaningful generic Q/DQ stage and + the full rewrite must complete before ``torch.export``. + + Args: + model: the ``nn.Module`` to quantize. + config: a ``coreai_opt.quantization.QuantizerConfig``. Defaults to + ``coreai_opt``'s default (int8 weight + activation) when ``None``. + """ + + def __init__(self, model: nn.Module, config: Optional[Any] = None) -> None: + from coreai_opt.quantization import Quantizer, QuantizerConfig + + self._quantizer = Quantizer(model, config or QuantizerConfig()) + self._prepared: Optional[fx.GraphModule] = None + + def prepare( + self, + example_inputs: Sequence[Any], + dynamic_shapes: Any = None, + ) -> fx.GraphModule: + """PT2E prepare: export + insert fake-quant. Returns the prepared graph.""" + self._prepared = self._quantizer.prepare( + tuple(example_inputs), dynamic_shapes=dynamic_shapes + ) + return self._prepared + + def calibration_mode(self): + """Context manager to collect activation statistics on calibration data.""" + return self._quantizer.calibration_mode() + + def training_mode(self): + """Context manager for quantization-aware training.""" + return self._quantizer.training_mode() + + def convert(self) -> fx.GraphModule: + """Produce the export-ready, fully quantized Core AI graph. + + Runs the entire ``coreai_opt`` ``finalize(CoreAI)`` (``convert_pt2e`` + + conv/bn fold + Q/DQ -> ``coreai::`` rewrite + kv-cache relocation). The + result has no fake-quant, so it can be strict-``torch.export``ed and then + lowered via ``to_edge_transform_and_lower([CoreAIPartitioner()])``. + """ + from coreai_opt.common import ExportBackend + + if self._prepared is None: + raise RuntimeError("Call prepare() before convert().") + return self._quantizer.finalize(backend=ExportBackend.CoreAI) diff --git a/backends/apple/coreai/test/test_quantizer.py b/backends/apple/coreai/test/test_quantizer.py new file mode 100644 index 00000000000..2271d0e39d6 --- /dev/null +++ b/backends/apple/coreai/test/test_quantizer.py @@ -0,0 +1,183 @@ +# 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. + +import unittest +from collections import Counter + +import torch +import torch.nn as nn + +from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner +from executorch.backends.apple.coreai.quantizer.quantizer import CoreAIQuantizer +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import ( + executorch_call_delegate, + get_lowered_backend_modules, +) + + +def _coreai_quant_op_names(graph_module): + """Names of call_function targets in a converted graph (test helper).""" + return [ + getattr(node.target, "__name__", str(node.target)) + for node in graph_module.graph.nodes + if node.op == "call_function" + ] + + +def _model(): + return nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 32)).eval() + + +class _ConvBN(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(3, 8, 3) + self.bn = nn.BatchNorm2d(8) + self.relu = nn.ReLU() + + def forward(self, x): + return self.relu(self.bn(self.conv(x))) + + +def _full_finalize(model, example_inputs): + """coreai_opt standalone finalize(CoreAI) for equivalence checks.""" + from coreai_opt.common import ExportBackend + from coreai_opt.quantization import Quantizer, QuantizerConfig + + q = Quantizer(model, QuantizerConfig()) + prepared = q.prepare(example_inputs) + with q.calibration_mode(): + prepared(*example_inputs) + return q.finalize(backend=ExportBackend.CoreAI) + + +class CoreAIQuantizerTest(unittest.TestCase): + def setUp(self): + self.example_inputs = (torch.randn(2, 32),) + + def _convert(self, model): + q = CoreAIQuantizer(model) + prepared = q.prepare(self.example_inputs) + with q.calibration_mode(): + prepared(*self.example_inputs) + return q.convert() + + def test_convert_produces_coreai_quant_ops(self): + names = set(_coreai_quant_op_names(self._convert(_model()))) + self.assertIn("quantize", names) + self.assertIn("dequantize", names) + self.assertIn("constexpr_blockwise_shift_scale", names) + + def test_convert_matches_full_finalize(self): + # convert() IS the full coreai_opt finalize(CoreAI). + converted = self._convert(_model()) + full = _full_finalize(_model(), self.example_inputs) + self.assertEqual( + Counter(_coreai_quant_op_names(converted)), + Counter(_coreai_quant_op_names(full)), + ) + + def test_convert_output_is_exportable(self): + # The key property: no fake-quant remains, so strict torch.export works. + converted = self._convert(_model()) + ep = torch.export.export(converted, self.example_inputs) + namespaces = { + getattr(n.target, "namespace", None) + for n in ep.graph.nodes + if n.op == "call_function" + } + # Quant ops are lowered into the coreai namespace after export. + self.assertIn("coreai", namespaces) + + def test_prepare_is_required_before_convert(self): + q = CoreAIQuantizer(_model()) + with self.assertRaises(RuntimeError): + q.convert() + + +class CoreAIQuantizerConvBNTest(unittest.TestCase): + def setUp(self): + self.example_inputs = (torch.randn(1, 3, 16, 16),) + + def _convert(self, model): + q = CoreAIQuantizer(model) + prepared = q.prepare(self.example_inputs) + with q.calibration_mode(): + prepared(*self.example_inputs) + return q.convert() + + def test_convert_folds_conv_bn(self): + # convert() folds conv+bn (pre-export) so no batch_norm survives. + names = _coreai_quant_op_names(self._convert(_ConvBN().eval())) + self.assertTrue(any("conv" in n for n in names), names) + self.assertFalse( + any("batch_norm" in n for n in names), + f"batch_norm was not folded: {names}", + ) + + def test_convert_matches_full_finalize_conv_bn(self): + converted = self._convert(_ConvBN().eval()) + full = _full_finalize(_ConvBN().eval(), self.example_inputs) + self.assertEqual( + Counter(_coreai_quant_op_names(converted)), + Counter(_coreai_quant_op_names(full)), + ) + + +class CoreAIQuantizedLoweringTest(unittest.TestCase): + """Full quantized lowering: quantize -> export -> to_edge_transform_and_lower.""" + + def setUp(self): + self.example_inputs = (torch.randn(2, 32),) + + def _lower(self): + q = CoreAIQuantizer(_model()) + prepared = q.prepare(self.example_inputs) + with q.calibration_mode(): + prepared(*self.example_inputs) + ep = torch.export.export(q.convert(), self.example_inputs) + return to_edge_transform_and_lower(ep, partitioner=[CoreAIPartitioner()]) + + def test_quantized_linear_lowers_to_coreai(self): + lowered = self._lower() + gm = lowered.exported_program().graph_module + + delegates = [ + n + for n in gm.graph.nodes + if n.op == "call_function" and n.target is executorch_call_delegate + ] + leftover = [ + n + for n in gm.graph.nodes + if n.op == "call_function" + and n.target is not executorch_call_delegate + and "getitem" not in str(n.target) + ] + # The quantized subgraph (incl. coreai:: quant ops) is delegated and + # converted by coreai-torch, leaving nothing outside the delegate. + self.assertGreaterEqual(len(delegates), 1) + self.assertEqual( + leftover, + [], + f"unexpected ops left outside the delegate: " + f"{[str(n.target) for n in leftover]}", + ) + + def test_quantized_lowers_through_to_executorch(self): + lowered = self._lower() + # A single CoreAI delegate, with its asset embedded (NDS present). + lbms = get_lowered_backend_modules(lowered.exported_program().graph_module) + self.assertEqual([lbm.backend_id for lbm in lbms], ["CoreAIBackend"]) + self.assertIsNotNone(lbms[0].named_data_store_output) + # It lowers all the way to a non-empty .pte. + buffer = bytes(lowered.to_executorch().buffer) + self.assertGreater(len(buffer), 0) + + +if __name__ == "__main__": + unittest.main()