From df6bcad622385b7f9f6c436e6aee37581495021f Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 28 Oct 2026 12:47:52 -0700 Subject: [PATCH] Lower quantized multiply through the Espressif operator library The opt-in export path now replaces supported quantized multiply patterns with direct esp_nn::quantized_mul_int8 calls. AtenToEspPass dispatches to per-operator functions that read Q/DQ arguments and combine folding with replacement. The initial contract is static, equal-shaped, contiguous int8 tensors with constant full-range per-tensor qparams and safe ESP-NN shift and overflow bounds. Replace the output quantizer's uses so existing output dequantizers and chained operations consume the direct result. Retain the original ATen operation and input dequantizers when a direct float consumer needs the pre-quantization result. Remove only unused nodes. Unsupported patterns, including multiple output quantizers, keep executable Q/DQ and portable compute. Also retain that pattern when the full int8 range can overflow float32 dequantization or multiplication, accounting for float32 rounding of scales and dequantized bounds. Add functional and out schemas, a fake implementation, and a Python integer reference. Share the delegate's Q31 multiplier encoder under backends/espressif/utils, preserving its finite-positive-input rounding and shift range. The existing EspNNQuantizer serves both export paths. Native C++ kernels and runtime registration are the next implementation step; this change establishes Python execution and PTE serialization. Use named multiply cases and EspTester.test_dialect() for operator expectations and numerical checks. Pass-manager tests cover export and API contracts without depending on exact Q/DQ counts. Start review with passes/aten_to_esp_pass.py, then ops/operators.py and utils/quantization.py for the integer contract, followed by test/ops/test_mul.py and the pass-manager tests. All 96 focused multiply, quantization-parameter, and export tests passed. The overflow regression preserves the portable saturated result for extreme qparams; range cases cover each input, their product, zero points, and float32 rounding. Scoped lint and whitespace checks passed. Authored with AI assistance using Codex. --- backends/espressif/lowering.py | 3 +- backends/espressif/ops/__init__.py | 5 + backends/espressif/ops/operators.py | 50 +++ backends/espressif/passes/aten_to_esp_pass.py | 157 ++++++++ backends/espressif/passes/esp_pass_manager.py | 4 +- backends/espressif/serialization/xts3.py | 25 +- backends/espressif/test/ops/__init__.py | 5 + backends/espressif/test/ops/test_mul.py | 140 +++++++ backends/espressif/test/test_pass_manager.py | 55 +-- .../test/test_quantization_params.py | 353 ++++++++++++++++++ backends/espressif/test/tester.py | 35 ++ backends/espressif/utils/__init__.py | 5 + backends/espressif/utils/quantization.py | 41 ++ 13 files changed, 818 insertions(+), 60 deletions(-) create mode 100644 backends/espressif/ops/__init__.py create mode 100644 backends/espressif/ops/operators.py create mode 100644 backends/espressif/passes/aten_to_esp_pass.py create mode 100644 backends/espressif/test/ops/__init__.py create mode 100644 backends/espressif/test/ops/test_mul.py create mode 100644 backends/espressif/test/test_quantization_params.py create mode 100644 backends/espressif/utils/__init__.py create mode 100644 backends/espressif/utils/quantization.py diff --git a/backends/espressif/lowering.py b/backends/espressif/lowering.py index f4d93198366..9998b993a6d 100644 --- a/backends/espressif/lowering.py +++ b/backends/espressif/lowering.py @@ -59,7 +59,8 @@ def to_espnn_edge( the program for the portable CPU kernels. Returns (edge_manager, report). With use_operator_library=True, run EspPassManager without partitioning. - This route currently preserves portable operators and executable Q/DQ. + Supported patterns become direct ESP-NN operators; others retain portable + operators and executable Q/DQ. Its report counts esp_nn_ops and lists remaining cpu_ops by method, excluding Q/DQ and structural getitem nodes. """ diff --git a/backends/espressif/ops/__init__.py b/backends/espressif/ops/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/espressif/ops/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/espressif/ops/operators.py b/backends/espressif/ops/operators.py new file mode 100644 index 00000000000..f6547936bb8 --- /dev/null +++ b/backends/espressif/ops/operators.py @@ -0,0 +1,50 @@ +# 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 torch +from executorch.backends.espressif.utils.quantization import requantize +from torch.library import impl, Library, register_fake + +lib = Library("esp_nn", "DEF") +lib.define( + "quantized_mul_int8(Tensor input1, Tensor input2, int input1_offset, " + "int input2_offset, int output_offset, int output_multiplier, int output_shift) -> Tensor" +) +lib.define( + "quantized_mul_int8.out(Tensor input1, Tensor input2, int input1_offset, " + "int input2_offset, int output_offset, int output_multiplier, int output_shift, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("esp_nn::quantized_mul_int8") +def quantized_mul_int8_meta( + input1, + input2, + input1_offset, + input2_offset, + output_offset, + output_multiplier, + output_shift, +): + return torch.empty_like(input1) + + +@impl(lib, "quantized_mul_int8", "CompositeExplicitAutograd") +def quantized_mul_int8( + input1, + input2, + input1_offset, + input2_offset, + output_offset, + output_multiplier, + output_shift, +): + product = (input1.to(torch.int32) + input1_offset) * ( + input2.to(torch.int32) + input2_offset + ) + result = requantize(product, output_multiplier, output_shift) + output_offset + return result.clamp(-128, 127).to(torch.int8) diff --git a/backends/espressif/passes/aten_to_esp_pass.py b/backends/espressif/passes/aten_to_esp_pass.py new file mode 100644 index 00000000000..73525236170 --- /dev/null +++ b/backends/espressif/passes/aten_to_esp_pass.py @@ -0,0 +1,157 @@ +# 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 math +from ctypes import c_float +from dataclasses import dataclass + +import executorch.backends.espressif.ops.operators # noqa: F401 +import torch +from executorch.backends.espressif.utils.quantization import quantize_multiplier +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_base import PassBase, PassResult + +_QUANTIZE = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +_DEQUANTIZE = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + + +@dataclass(frozen=True) +class _QParams: + scale: float + zp: int + + +def _qparams(node: Node) -> _QParams | None: + if len(node.args) < 6: + return None + scale, zp, qmin, qmax, dtype = node.args[1:6] + if ( + dtype != torch.int8 + or any(type(value) is not int for value in (zp, qmin, qmax)) + or type(scale) not in (float, int) + ): + return None + if not math.isfinite(scale) or scale <= 0: + raise ValueError(f"{node.name}: quantization scale must be finite and positive") + if not -128 <= qmin <= zp <= qmax <= 127: + raise ValueError(f"{node.name}: invalid int8 quantization range or zero point") + if (qmin, qmax) != (-128, 127): + return None + return _QParams(float(scale), zp) + + +def _mul_multiplier( + input1: _QParams, input2: _QParams, output: _QParams +) -> tuple[int, int] | None: + rescale = input1.scale * input2.scale / output.scale + if not math.isfinite(rescale) or rescale <= 0: + return None + try: + multiplier, shift = quantize_multiplier(rescale) + except ValueError: + return None + # ESP-NN forms (1 << -shift) - 1 with signed int32 arithmetic. + if shift < -30: + return None + max_input1 = max(abs(-128 - input1.zp), abs(127 - input1.zp)) + max_input2 = max(abs(-128 - input2.zp), abs(127 - input2.zp)) + max_product = max_input1 * max_input2 + if max_product << max(shift, 0) > torch.iinfo(torch.int32).max - abs(output.zp): + return None + # Account for float32 rounding of both the scales and dequantized inputs. + max_float1 = c_float(max_input1 * c_float(input1.scale).value).value + max_float2 = c_float(max_input2 * c_float(input2.scale).value).value + if ( + not math.isfinite(max_float1) + or not math.isfinite(max_float2) + or max_float1 * max_float2 > torch.finfo(torch.float32).max + ): + return None + return multiplier, shift + + +def _mul_args(node: Node, quantize: Node) -> tuple | None: + inputs = node.args[:2] + if len(inputs) != 2 or not all( + isinstance(arg, Node) and arg.target == _DEQUANTIZE for arg in inputs + ): + return None + params = [_qparams(arg) for arg in (*inputs, quantize)] + if any(param is None for param in params): + return None + values = [arg.args[0].meta.get("val") for arg in inputs] + [ + quantize.meta.get("val") + ] + if not all( + isinstance(value, torch.Tensor) + and value.dtype == torch.int8 + and all(type(dim) is int for dim in value.shape) + and 0 < value.numel() <= torch.iinfo(torch.int32).max + and value.is_contiguous() + for value in values + ) or not all(value.shape == values[0].shape for value in values[1:]): + return None + if not all( + isinstance(arg.meta.get("val"), torch.Tensor) + and arg.meta["val"].dtype == torch.float32 + for arg in (*inputs, node) + ): + return None + rescale = _mul_multiplier(*params) + if rescale is None: + return None + input1, input2, output = params + return ( + inputs[0].args[0], + inputs[1].args[0], + -input1.zp, + -input2.zp, + output.zp, + *rescale, + ) + + +def _lower_mul(node: Node) -> bool: + quantizers = [ + user for user in node.users if user.target == _QUANTIZE and user.args[0] is node + ] + if len(quantizers) != 1: + return False + quantize = quantizers[0] + args = _mul_args(node, quantize) + if args is None: + return False + with node.graph.inserting_before(quantize): + replacement = node.graph.call_function( + exir_ops.edge.esp_nn.quantized_mul_int8.default, args + ) + replacement.meta = quantize.meta.copy() + # Replace the quantized branch; direct float consumers still need the ATen op. + quantize.replace_all_uses_with(replacement) + node.graph.erase_node(quantize) + return True + + +_LOWERINGS = {exir_ops.edge.aten.mul.Tensor: _lower_mul} + + +class AtenToEspPass(PassBase): + """Fold supported Q/DQ patterns directly into ESP-NN operator calls.""" + + def call(self, graph_module: GraphModule) -> PassResult: + modified = False + for node in list(graph_module.graph.nodes): + if node.op != "call_function": + continue + lowering = _LOWERINGS.get(node.target) + if lowering is not None: + modified |= lowering(node) + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/espressif/passes/esp_pass_manager.py b/backends/espressif/passes/esp_pass_manager.py index 1c77850af81..d122ab8218f 100644 --- a/backends/espressif/passes/esp_pass_manager.py +++ b/backends/espressif/passes/esp_pass_manager.py @@ -6,9 +6,11 @@ from executorch.exir.pass_manager import ExportedProgramPassManager +from .aten_to_esp_pass import AtenToEspPass + class EspPassManager(ExportedProgramPassManager): """Run Espressif operator-library transformations on an Edge program.""" def __init__(self) -> None: - super().__init__(passes=[]) + super().__init__(passes=[AtenToEspPass()]) diff --git a/backends/espressif/serialization/xts3.py b/backends/espressif/serialization/xts3.py index 0a99476b05e..7ec71b95a1d 100644 --- a/backends/espressif/serialization/xts3.py +++ b/backends/espressif/serialization/xts3.py @@ -11,10 +11,12 @@ import struct from dataclasses import dataclass, field from enum import IntEnum -from math import frexp, prod +from math import prod import torch - +from executorch.backends.espressif.utils.quantization import ( + quantize_multiplier as _quantize_multiplier, +) MAGIC = b"XTS3" XTS3_VERSION = 6 @@ -829,25 +831,6 @@ def _is_relu6(node) -> bool: return minimum == 0.0 and maximum == 6.0 -def _quantize_multiplier(real_multiplier: float) -> tuple[int, int]: - if not real_multiplier > 0.0: - raise ValueError( - f"Quantized multiplier must be positive, got {real_multiplier}" - ) - significand, exponent = frexp(real_multiplier) - multiplier = round(significand * (1 << 31)) - if multiplier == (1 << 31): - multiplier >>= 1 - exponent += 1 - if multiplier > 0x7FFFFFFF: - multiplier = 0x7FFFFFFF - if exponent < -31 or exponent > 30: - raise ValueError( - f"Requantization multiplier {real_multiplier} is outside the ESP-NN shift range" - ) - return int(multiplier), int(exponent) - - def _quantize_multiplier_smaller_than_one(real_multiplier: float) -> tuple[int, int]: if not 0.0 < real_multiplier < 1.0: raise ValueError(f"Expected a multiplier in (0, 1), got {real_multiplier}") diff --git a/backends/espressif/test/ops/__init__.py b/backends/espressif/test/ops/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/espressif/test/ops/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/espressif/test/ops/test_mul.py b/backends/espressif/test/ops/test_mul.py new file mode 100644 index 00000000000..84598bf989d --- /dev/null +++ b/backends/espressif/test/ops/test_mul.py @@ -0,0 +1,140 @@ +# 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 pytest +import torch +from executorch.backends.espressif.test.tester import ( + EspTestCase, + EspTester, + ramp_tensor, +) +from executorch.backends.test.suite.operators.test_mul import Model +from executorch.exir._serialize._program import deserialize_pte_binary +from executorch.exir.dialects._ops import ops as exir_ops + +_MUL = exir_ops.edge.aten.mul.Tensor +_ESP_MUL = exir_ops.edge.esp_nn.quantized_mul_int8.default + + +class EspMul(Model): + ops_before_transforms = {_MUL: 1} + ops_after_transforms = {_MUL: 0, _ESP_MUL: 1} + + +class EspSelfMul(EspMul): + def forward(self, x): + return x * x + + +class EspConstantMul(EspMul): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.linspace(3, -2, 17)) + + def forward(self, x): + return x * self.weight + + +class EspMulChain(EspMul): + ops_before_transforms = {_MUL: 2} + ops_after_transforms = {_MUL: 0, _ESP_MUL: 2} + + def forward(self, x, y): + product = x * y + return product * x, torch.sin(product), torch.sin(x) + + +class EspMulWithFloatUsers(EspMul): + def forward(self, x, y): + product = x * y + return product, torch.sin(x), torch.sin(product) + + +class EspBroadcastMul(EspMulWithFloatUsers): + ops_after_transforms = {_MUL: 1, _ESP_MUL: 0} + + +class EspMulAdd(EspMul): + ops_after_transforms = { + _MUL: 0, + _ESP_MUL: 1, + exir_ops.edge.aten.add.Tensor: 1, + } + + def forward(self, x, y): + return x * y + y + + +test_cases = { + "scalar": EspTestCase( + EspMul(), + (ramp_tensor(-1.5, 2.5, ()), ramp_tensor(3, -4, ())), + ), + "tensor_rank_1": EspTestCase( + EspMul(), + (ramp_tensor(-1.5, 2.5, (17,)), ramp_tensor(3, -4, (17,))), + ), + "tensor_rank_3": EspTestCase( + EspMul(), + (ramp_tensor(-1.5, 2.5, (2, 3, 5)), ramp_tensor(3, -4, (2, 3, 5))), + ), + "tensor_rank_4": EspTestCase( + EspMul(), + (ramp_tensor(-1.5, 2.5, (1, 2, 3, 4)), ramp_tensor(3, -4, (1, 2, 3, 4))), + ), + "self": EspTestCase(EspSelfMul(), (ramp_tensor(-1, 2, (17,)),)), + "constant": EspTestCase(EspConstantMul(), (ramp_tensor(-1, 2, (17,)),)), + "chain": EspTestCase( + EspMulChain(), + (ramp_tensor(-1, 1, (17,)), ramp_tensor(2, -2, (17,))), + ), + "float_users": EspTestCase( + EspMulWithFloatUsers(), + (ramp_tensor(-1, 1, (2, 3)), ramp_tensor(-2, 2, (2, 3))), + ), + "broadcast": EspTestCase( + EspBroadcastMul(), + (ramp_tensor(-1, 1, (2, 3)), ramp_tensor(-2, 2, (3,))), + ), +} + + +@pytest.mark.parametrize("test_case", test_cases.values(), ids=test_cases.keys()) +def test_dialect_mul(test_case): + tester = EspTester(test_case.model, test_case.example_inputs) + tester.test_dialect( + test_case.model.ops_before_transforms, + test_case.model.ops_after_transforms, + qtol=1, + atol=0, + rtol=torch.finfo(torch.float32).eps, + ) + tester.to_executorch().serialize() + program = deserialize_pte_binary(tester.get_artifact()).program + assert not program.execution_plan[0].delegates + operators = {(op.name, op.overload) for op in program.execution_plan[0].operators} + assert (("esp_nn::quantized_mul_int8", "out") in operators) == ( + test_case.model.ops_after_transforms[_ESP_MUL] > 0 + ) + + +def test_mul_keeps_dequantization_for_portable_add(): + model = EspMulAdd() + inputs = (ramp_tensor(-1, 1, (17,)), ramp_tensor(2, -2, (17,))) + tester = EspTester(model, inputs) + tester.test_dialect( + model.ops_before_transforms, model.ops_after_transforms, qtol=1, atol=0, rtol=0 + ) + add = next( + node + for node in tester.get_artifact().exported_program().graph.nodes + if node.target == exir_ops.edge.aten.add.Tensor + ) + assert all( + arg.target == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + for arg in add.args[:2] + ) + tester.to_executorch().serialize() diff --git a/backends/espressif/test/test_pass_manager.py b/backends/espressif/test/test_pass_manager.py index 755cc9b9247..1fbb32dd5fd 100644 --- a/backends/espressif/test/test_pass_manager.py +++ b/backends/espressif/test/test_pass_manager.py @@ -11,7 +11,6 @@ from executorch.backends.espressif.test.tester import EspRunPasses, EspTester from executorch.exir import to_edge, to_edge_transform_and_lower from executorch.exir._serialize._program import deserialize_pte_binary -from executorch.exir.dialects._ops import ops as exir_ops class Multiply(torch.nn.Module): @@ -25,28 +24,23 @@ def forward(self, x, y): return product, torch.sin(x), torch.sin(product) +def _output_scale(program): + output = next(node for node in program.graph.nodes if node.op == "output") + return output.args[0][0].args[1] + + @pytest.mark.parametrize("entry_point", ["separate", "combined"]) -@pytest.mark.parametrize("broadcast", [False, True]) -def test_quantized_and_float_operations_survive_export(entry_point, broadcast): +def test_entry_points_preserve_program_behavior(entry_point): inputs = ( torch.linspace(-1, 1, 6).reshape(2, 3), - torch.linspace(-2, 2, 3 if broadcast else 6).reshape( - (3,) if broadcast else (2, 3) - ), + torch.linspace(-2, 2, 6).reshape(2, 3), ) tester = EspTester(MultiplyAndFloatUsers(), inputs).quantize().export() if entry_point == "combined": tester.to_edge_transform_and_lower() else: tester.to_edge().run_passes() - tester.check_node_count( - { - exir_ops.edge.aten.mul.Tensor: 1, - exir_ops.edge.aten.sin.default: 2, - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default: 3, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default: 3, - } - ).run_method_and_compare_outputs(inputs=inputs, atol=0, rtol=0) + tester.run_method_and_compare_outputs(inputs=inputs, atol=0, rtol=0, qtol=1) edge = tester.get_artifact() assert edge.compile_config._check_ir_validity edge.exported_program().validate() @@ -54,12 +48,6 @@ def test_quantized_and_float_operations_survive_export(entry_point, broadcast): program = deserialize_pte_binary(tester.get_artifact()).program assert program.execution_plan[0].name == "forward" assert not program.execution_plan[0].delegates - assert {op.name for op in program.execution_plan[0].operators} >= { - "aten::mul", - "aten::sin", - "quantized_decomposed::quantize_per_tensor", - "quantized_decomposed::dequantize_per_tensor", - } def test_public_entry_point_keeps_delegate_default(): @@ -72,16 +60,13 @@ def test_public_entry_point_keeps_delegate_default(): assert not delegate_report["cpu_ops"] assert delegated.to_executorch().executorch_program.execution_plan[0].delegates - edge, report = to_espnn_edge(program, use_operator_library=True) - assert report["esp_nn_ops"] == 0 - assert [(op["method"], op["target"]) for op in report["cpu_ops"]] == [ - ("forward", "aten::mul.Tensor") - ] + edge, report = to_espnn_edge(program, False, use_operator_library=True) + assert report == {"esp_nn_ops": 1, "cpu_ops": []} assert edge.compile_config._check_ir_validity torch.testing.assert_close( edge.exported_program().module()(*inputs), program.module()(*inputs), - atol=0, + atol=_output_scale(program), rtol=0, ) serialized = deserialize_pte_binary(edge.to_executorch().buffer).program @@ -90,8 +75,10 @@ def test_public_entry_point_keeps_delegate_default(): def test_operator_library_strict_mode_rejects_remaining_compute(): inputs = (torch.linspace(-1, 1, 6), torch.linspace(-2, 2, 6)) - program = EspTester(Multiply(), inputs).quantize().export().get_artifact() - with pytest.raises(RuntimeError, match=r"forward\.\w+ \(aten::mul\.Tensor\)"): + program = ( + EspTester(MultiplyAndFloatUsers(), inputs).quantize().export().get_artifact() + ) + with pytest.raises(RuntimeError, match=r"forward\.\w+ \(aten::sin\)"): to_espnn_edge(program, False, use_operator_library=True) @@ -103,14 +90,11 @@ def test_operator_library_strict_mode_accepts_no_compute(): def test_operator_library_uses_standard_edge_decompositions(): inputs = (torch.linspace(-1, 1, 6).reshape(2, 3),) - program = ( - EspTester(torch.nn.Linear(3, 2), inputs).quantize().export().get_artifact() - ) + program = torch.export.export(torch.nn.Linear(3, 2), inputs) edge, report = to_espnn_edge(program, use_operator_library=True) targets = {op["target"] for op in report["cpu_ops"]} assert "aten::linear" not in targets assert "aten::addmm" in targets - assert not any(target.startswith("quantized_decomposed::") for target in targets) torch.testing.assert_close( edge.exported_program().module()(*inputs), program.module()(*inputs) ) @@ -169,15 +153,12 @@ def test_operator_library_reports_all_methods(): for name in ("forward", "other") } edge, report = to_espnn_edge(programs, use_operator_library=True) - assert report["esp_nn_ops"] == 0 - assert [(op["method"], op["target"]) for op in report["cpu_ops"]] == [ - (name, "aten::mul.Tensor") for name in ("forward", "other") - ] + assert report == {"esp_nn_ops": 2, "cpu_ops": []} for name, program in programs.items(): torch.testing.assert_close( edge.exported_program(name).module()(*inputs), program.module()(*inputs), - atol=0, + atol=_output_scale(program), rtol=0, ) serialized = deserialize_pte_binary(edge.to_executorch().buffer).program diff --git a/backends/espressif/test/test_quantization_params.py b/backends/espressif/test/test_quantization_params.py new file mode 100644 index 00000000000..f3fd353b5f4 --- /dev/null +++ b/backends/espressif/test/test_quantization_params.py @@ -0,0 +1,353 @@ +# 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 copy + +import pytest +import torch +from executorch.backends.espressif.passes import EspPassManager +from executorch.backends.espressif.passes.aten_to_esp_pass import AtenToEspPass +from executorch.backends.espressif.reference.esp_nn_reference import ( + multiply_by_quantized_multiplier, +) +from executorch.backends.espressif.serialization.xts3 import _quantize_multiplier +from executorch.backends.espressif.test.tester import EspTester +from executorch.backends.espressif.utils.quantization import ( + quantize_multiplier, + requantize, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops + +_QUANTIZE = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +_MUL = exir_ops.edge.aten.mul.Tensor +_ESP_MUL = exir_ops.edge.esp_nn.quantized_mul_int8.default + + +class Multiply(torch.nn.Module): + def forward(self, x, y): + return x * y + + +@pytest.fixture(scope="module") +def exported_multiply(): + inputs = (torch.linspace(-1, 1, 6).reshape(2, 3), torch.ones(2, 3)) + return ( + EspTester(Multiply(), inputs) + .quantize() + .export() + .to_edge() + .get_artifact() + .exported_program() + ) + + +@pytest.fixture +def multiply_program(exported_multiply): + return copy.deepcopy(exported_multiply) + + +def _pattern(graph_module): + mul = next(node for node in graph_module.graph.nodes if node.target == _MUL) + quantize = next(node for node in mul.users if node.target == _QUANTIZE) + return *mul.args, quantize + + +@pytest.mark.parametrize( + "scale, expected", + [ + (2**-32, (1 << 30, -31)), + (2**-31, (1 << 30, -30)), + (0.25, (1 << 30, -1)), + (0.5, (1 << 30, 0)), + (1.0, (1 << 30, 1)), + (2**29, (1 << 30, 30)), + ((2**30 + 0.5) / 2**31, (1 << 30, 0)), + ((2**30 + 1.5) / 2**31, ((1 << 30) + 2, 0)), + (1 - 2**-32, (1 << 30, 1)), + ], +) +def test_multiplier_encoding_preserves_delegate_rounding(scale, expected): + assert quantize_multiplier(scale) == expected + assert _quantize_multiplier(scale) == expected + + +@pytest.mark.parametrize("scale", [0, -1, float("nan"), float("inf"), -float("inf")]) +def test_multiplier_requires_finite_positive_scale(scale): + with pytest.raises(ValueError, match="finite and positive"): + quantize_multiplier(scale) + + +@pytest.mark.parametrize("scale", [2**-33, 2**30]) +def test_multiplier_shift_range(scale): + with pytest.raises(ValueError, match="shift range"): + quantize_multiplier(scale) + + +def test_folding_and_replacement(multiply_program): + graph = multiply_program.graph_module + input1, input2, output = _pattern(graph) + for node, scale, zero_point in ( + (input1, 0.25, -17), + (input2, 0.125, 23), + (output, 0.5, -7), + ): + node.args = (node.args[0], scale, zero_point, -128, 127, torch.int8) + graph.recompile() + inputs = (torch.linspace(-1, 1, 6).reshape(2, 3), torch.ones(2, 3)) + expected = graph(*inputs) + lowered = AtenToEspPass()(graph).graph_module + mul = next(node for node in lowered.graph.nodes if node.target == _ESP_MUL) + assert mul.args[2:] == (17, -23, -7, 1 << 30, -3) + assert all(arg.meta["val"].dtype == torch.int8 for arg in mul.args[:2]) + assert output not in lowered.graph.nodes + assert not any(node.target == _MUL for node in lowered.graph.nodes) + torch.testing.assert_close(lowered(*inputs), expected, atol=0.5, rtol=0) + + +@pytest.mark.parametrize( + "rescale, input_zero_point, expected_shift", + [ + (2**-32, 0, None), + (2**-31, 0, -30), + (2**15, 0, 16), + (2**16, 0, None), + (2**14, -128, 15), + (2**15, -128, None), + (2**29, 0, None), + (2**30, 0, None), + ], +) +def test_kernel_shift_and_overflow_limits( + multiply_program, rescale, input_zero_point, expected_shift +): + input1, input2, output = _pattern(multiply_program.graph_module) + for node in (input1, input2): + node.args = (node.args[0], 1.0, input_zero_point, -128, 127, torch.int8) + output.args = (output.args[0], 1 / rescale, 127, -128, 127, torch.int8) + result = EspPassManager()(multiply_program) + mul = next( + ( + node + for node in result.exported_program.graph.nodes + if node.target == _ESP_MUL + ), + None, + ) + if expected_shift is None: + assert mul is None + else: + assert mul.args[-1] == expected_shift + assert mul.args[4] == 127 + + +@pytest.mark.parametrize("input_scale", [1e-200, 1e200]) +def test_unrepresentable_rescale_stays_portable(multiply_program, input_scale): + input1, input2, output = _pattern(multiply_program.graph_module) + for node in (input1, input2): + node.args = (node.args[0], input_scale, 0, -128, 127, torch.int8) + output.args = (output.args[0], 1.0, 0, -128, 127, torch.int8) + assert not AtenToEspPass()(multiply_program.graph_module).modified + assert output in multiply_program.graph.nodes + + +@pytest.mark.parametrize( + "scales, zero_point, supported", + [ + ((1e38, 1e-38, 1.0), 0, False), + ((1e-38, 1e38, 1.0), 0, False), + ((1e19, 1e19, 1e38), 0, False), + ((1e17, 1e17, 1e34), 0, True), + ((1e17, 1e17, 1e34), -128, False), + ( + (1.2e17, torch.finfo(torch.float32).max / (128 * 128 * 1.2e17), 1e34), + 0, + False, + ), + ], + ids=["input1", "input2", "product", "finite", "offset", "float32_rounding"], +) +def test_float32_range_limits(multiply_program, scales, zero_point, supported): + graph = multiply_program.graph_module + for node, scale in zip(_pattern(graph), scales): + node.args = (node.args[0], scale, zero_point, -128, 127, torch.int8) + assert AtenToEspPass()(graph).modified == supported + + +def test_float32_overflow_preserves_portable_results(multiply_program): + graph = multiply_program.graph_module + input1, input2, output = _pattern(graph) + for dequantize in (input1, input2): + quantize = dequantize.args[0] + for node in (quantize, dequantize): + node.args = (node.args[0], 1e19, 0, -128, 127, torch.int8) + output.args = (output.args[0], 1e38, 0, -128, 127, torch.int8) + graph_output = next(node for node in graph.graph.nodes if node.op == "output") + graph_output.args = ((output,),) + graph.graph.eliminate_dead_code() + graph.recompile() + inputs = torch.tensor([[1.0, 2.0, 10.0], [1.0, 2.0, 10.0]]) * 1e19 + expected = torch.tensor([[1, 127, 127], [1, 127, 127]], dtype=torch.int8) + torch.testing.assert_close(graph(inputs, inputs)[0], expected, atol=0, rtol=0) + + assert not AtenToEspPass()(graph).modified + assert output in graph.graph.nodes + torch.testing.assert_close(graph(inputs, inputs)[0], expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("operand", [0, 1, 2]) +@pytest.mark.parametrize( + "qparams, message", + [ + ((0.0, 0, -128, 127), "finite and positive"), + ((-1.0, 0, -128, 127), "finite and positive"), + ((float("nan"), 0, -128, 127), "finite and positive"), + ((float("inf"), 0, -128, 127), "finite and positive"), + ((1.0, -129, -128, 127), "range or zero point"), + ((1.0, 128, -128, 127), "range or zero point"), + ((1.0, 0, -129, 127), "range or zero point"), + ((1.0, 0, 1, -1), "range or zero point"), + ], +) +def test_malformed_qparams_fail_with_node_name( + multiply_program, operand, qparams, message +): + node = _pattern(multiply_program.graph_module)[operand] + node.args = (node.args[0], *qparams, torch.int8) + with pytest.raises(ValueError, match=f"{node.name}:.*{message}"): + AtenToEspPass()(multiply_program.graph_module) + + +@pytest.mark.parametrize( + "value", + [ + torch.empty(3, 2, dtype=torch.int8, device="meta"), + torch.empty(3, 2, dtype=torch.int8, device="meta").t(), + torch.empty(0, dtype=torch.int8, device="meta"), + torch.empty(1 << 31, dtype=torch.int8, device="meta"), + ], + ids=["broadcast", "noncontiguous", "empty", "too_many_elements"], +) +def test_unsupported_storage_stays_portable(multiply_program, value): + input1, _, output = _pattern(multiply_program.graph_module) + input1.args[0].meta["val"] = value + assert not AtenToEspPass()(multiply_program.graph_module).modified + assert output in multiply_program.graph.nodes + + +@pytest.mark.parametrize( + "dtype, quant_min, quant_max", [(torch.uint8, 0, 255), (torch.int8, -127, 127)] +) +def test_unsupported_quantization_stays_portable( + multiply_program, dtype, quant_min, quant_max +): + input1, _, output = _pattern(multiply_program.graph_module) + input1.args = (input1.args[0], 1.0, 0, quant_min, quant_max, dtype) + assert not AtenToEspPass()(multiply_program.graph_module).modified + assert output in multiply_program.graph.nodes + + +def test_dynamic_shapes_stay_portable(): + inputs = (torch.ones(2, 3), torch.ones(2, 3)) + quantized = EspTester(Multiply(), inputs).quantize().get_artifact() + batch = torch.export.Dim("batch", min=1, max=8) + program = torch.export.export( + quantized, inputs, dynamic_shapes=({0: batch}, {0: batch}) + ) + graph = to_edge(program).exported_program().graph_module + _, _, output = _pattern(graph) + assert not AtenToEspPass()(graph).modified + assert output in graph.graph.nodes + + +def test_multiple_output_quantizers_stay_portable(multiply_program): + graph = multiply_program.graph_module + _, _, output = _pattern(graph) + mul = output.args[0] + with graph.graph.inserting_after(output): + other = graph.graph.call_function( + _QUANTIZE, (mul, 2.0, 13, -128, 127, torch.int8) + ) + other.meta = output.meta.copy() + graph_output = next(node for node in graph.graph.nodes if node.op == "output") + graph_output.args = ((*graph_output.args[0], other),) + graph.recompile() + inputs = (torch.linspace(-1, 1, 6).reshape(2, 3), torch.ones(2, 3)) + expected = graph(*inputs) + assert not AtenToEspPass()(graph).modified + torch.testing.assert_close(graph(*inputs), expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("float_user_first", [False, True]) +def test_direct_float_branch_is_preserved(float_user_first): + class BranchedMultiply(torch.nn.Module): + def forward(self, x, y): + product = x * y + return product, torch.sin(product) + + inputs = ( + torch.linspace(-1, 1, 6).reshape(2, 3), + torch.linspace(-2, 2, 6).reshape(2, 3), + ) + tester = EspTester(BranchedMultiply(), inputs).quantize().export().to_edge() + graph = tester.get_artifact().exported_program().graph_module + input1, input2, quantize = _pattern(graph) + mul = quantize.args[0] + float_user = next( + node + for node in graph.graph.nodes + if node.target == exir_ops.edge.aten.sin.default + ) + # PT2E normally puts sin after output DQ; exercise a direct float consumer. + float_user.args = (mul,) + if float_user_first: + with graph.graph.inserting_after(quantize): + new_quantize = graph.graph.node_copy(quantize) + quantize.replace_all_uses_with(new_quantize) + graph.graph.erase_node(quantize) + quantize = new_quantize + graph.recompile() + expected = graph(*inputs) + assert not torch.equal(torch.sin(expected[0]), expected[1]) + + tester.run_passes() + program = tester.get_artifact().exported_program() + lowered = program.graph_module + esp_mul = next(node for node in lowered.graph.nodes if node.target == _ESP_MUL) + assert quantize not in lowered.graph.nodes + assert mul in lowered.graph.nodes + assert mul.args == (input1, input2) + assert float_user.args == (mul,) + output_dequantize = next(iter(esp_mul.users)) + assert output_dequantize.target == ( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + ) + actual = lowered(*inputs) + torch.testing.assert_close( + actual[0], expected[0], atol=output_dequantize.args[1], rtol=0 + ) + torch.testing.assert_close(actual[1], expected[1], atol=0, rtol=0) + assert not EspPassManager()(program).modified + tester.to_executorch().serialize() + + +def test_lowering_is_idempotent(multiply_program): + first = EspPassManager()(multiply_program) + assert first.modified + assert not EspPassManager()(first.exported_program).modified + + +@pytest.mark.parametrize("scale", [0.03, 0.25, 0.5, 0.7, 1.0, 1.3, 2.0, 32768.0]) +def test_requantization_matches_existing_integer_reference(scale): + values = torch.arange(-128, 128, dtype=torch.int64) + products = (values[:, None] * values[None, :]).flatten() + multiplier, shift = quantize_multiplier(scale) + expected = torch.from_numpy( + multiply_by_quantized_multiplier(products.numpy(), multiplier, shift) + ).to(torch.int32) + torch.testing.assert_close( + requantize(products, multiplier, shift), expected, atol=0, rtol=0 + ) diff --git a/backends/espressif/test/tester.py b/backends/espressif/test/tester.py index c45cc300672..75458cf3dc3 100644 --- a/backends/espressif/test/tester.py +++ b/backends/espressif/test/tester.py @@ -4,6 +4,10 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import math +from dataclasses import dataclass + +import torch from executorch.backends.espressif.passes import EspPassManager from executorch.backends.espressif.quantizer import EspNNQuantizer from executorch.backends.test.harness import Tester @@ -50,3 +54,34 @@ def __init__(self, module, example_inputs, calibration_samples=None): stages[StageType.RUN_PASSES] = EspRunPasses stages[StageType.TO_EDGE_TRANSFORM_AND_LOWER] = EspToEdgeTransformAndLower super().__init__(module, example_inputs, stages) + + def test_dialect( + self, + ops_before_transforms, + ops_after_transforms, + *, + qtol=0, + atol=1e-3, + rtol=1e-3, + ): + return ( + self.quantize() + .export() + .to_edge() + .check_node_count(ops_before_transforms) + .run_passes() + .check_node_count(ops_after_transforms) + .run_method_and_compare_outputs( + inputs=self.example_inputs, qtol=qtol, atol=atol, rtol=rtol + ) + ) + + +@dataclass +class EspTestCase: + model: torch.nn.Module + example_inputs: tuple[torch.Tensor, ...] + + +def ramp_tensor(start: float, end: float, shape: tuple[int, ...]) -> torch.Tensor: + return torch.linspace(start, end, math.prod(shape)).reshape(shape) diff --git a/backends/espressif/utils/__init__.py b/backends/espressif/utils/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/espressif/utils/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/espressif/utils/quantization.py b/backends/espressif/utils/quantization.py new file mode 100644 index 00000000000..81a3ec2a729 --- /dev/null +++ b/backends/espressif/utils/quantization.py @@ -0,0 +1,41 @@ +# Copyright 2026 Mwaniki Njagi. +# 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 math import frexp, isfinite + +import torch + + +def quantize_multiplier(real_multiplier: float) -> tuple[int, int]: + if not isfinite(real_multiplier) or real_multiplier <= 0.0: + raise ValueError( + f"Quantized multiplier must be finite and positive, got {real_multiplier}" + ) + significand, exponent = frexp(real_multiplier) + multiplier = round(significand * (1 << 31)) + if multiplier == (1 << 31): + multiplier >>= 1 + exponent += 1 + if exponent < -31 or exponent > 30: + raise ValueError( + f"Requantization multiplier {real_multiplier} is outside the ESP-NN shift range" + ) + return int(multiplier), int(exponent) + + +def requantize(value: torch.Tensor, multiplier: int, shift: int) -> torch.Tensor: + """ESP-NN's two-rounding Q31 rescale, for values whose left shift fits int32.""" + value = value.to(torch.int64) << max(shift, 0) + # High multiplication rounds ties upward; division rounds ties away from zero. + result = (value * multiplier + (1 << 30)) >> 31 + if shift < 0: + mask = (1 << -shift) - 1 + remainder = result & mask + result = result >> -shift + threshold = (mask >> 1) + (result < 0).to(torch.int64) + result = result + (remainder > threshold) + return result.to(torch.int32)