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)