From d9bf9c77507a128d96a57e5f7ee8a91fa085d3b6 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 28 Oct 2026 11:26:53 -0700 Subject: [PATCH] Add opt-in Espressif operator-library export pipeline Introduce EspPassManager and a frontend EspTester using EspNNQuantizer. The new keyword-only export option preserves portable compute and executable Q/DQ with Edge IR checks enabled, reports remaining compute, and enforces strict fallback handling. Existing callers continue to use the delegate path. Validated with the complete 67-test Espressif host suite, including 11 new frontend tests covering both export routes, shared float consumers, strict fallback, multiple methods, constants, and serialization. Scoped lintrunner and whitespace checks passed. Authored with AI assistance using Codex. --- backends/espressif/lowering.py | 61 +++++- backends/espressif/passes/__init__.py | 9 + backends/espressif/passes/esp_pass_manager.py | 14 ++ backends/espressif/test/test_pass_manager.py | 185 ++++++++++++++++++ backends/espressif/test/tester.py | 52 +++++ 5 files changed, 319 insertions(+), 2 deletions(-) create mode 100644 backends/espressif/passes/__init__.py create mode 100644 backends/espressif/passes/esp_pass_manager.py create mode 100644 backends/espressif/test/test_pass_manager.py create mode 100644 backends/espressif/test/tester.py diff --git a/backends/espressif/lowering.py b/backends/espressif/lowering.py index 67e4063b54d..f4d93198366 100644 --- a/backends/espressif/lowering.py +++ b/backends/espressif/lowering.py @@ -8,21 +8,78 @@ from __future__ import annotations +import operator + import executorch.kernels.quantized # noqa: F401 registers quantize/dequantize out variants from executorch.backends.espressif.partition.espnn_partitioner import EspNNPartitioner -from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from executorch.backends.espressif.passes import EspPassManager +from executorch.exir import ( + EdgeCompileConfig, + EdgeProgramManager, + to_edge_transform_and_lower, +) from executorch.exir.backend.canonical_partitioners.duplicate_dequant_node_pass import ( DuplicateDequantNodePass, ) +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch._ops import OpOverload + +def _operator_library_report(edge: EdgeProgramManager) -> dict: + esp_nn_ops = 0 + cpu_ops = [] + for method in sorted(edge.methods): + for node in edge.exported_program(method).graph.nodes: + if node.op != "call_function" or node.target == operator.getitem: + continue + target = node.target + if isinstance(target, (EdgeOpOverload, OpOverload)): + if target._schema.name in { + "quantized_decomposed::quantize_per_tensor", + "quantized_decomposed::dequantize_per_tensor", + "quantized_decomposed::quantize_per_channel", + "quantized_decomposed::dequantize_per_channel", + }: + continue + if target.namespace == "esp_nn": + esp_nn_ops += 1 + continue + target = target.name() + cpu_ops.append({"method": method, "name": node.name, "target": str(target)}) + return {"esp_nn_ops": esp_nn_ops, "cpu_ops": cpu_ops} -def to_espnn_edge(program, allow_cpu_fallback: bool = True): + +def to_espnn_edge( + program, allow_cpu_fallback: bool = True, *, use_operator_library: bool = False +): """ Lower a PT2E-quantized torch.export program for the ESP32 ExecuTorch runtime: int8 regions become ESP-NN delegates, anything else stays in 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. + Its report counts esp_nn_ops and lists remaining cpu_ops by method, + excluding Q/DQ and structural getitem nodes. """ + if use_operator_library: + edge = to_edge_transform_and_lower( + program, + transform_passes=EspPassManager(), + compile_config=EdgeCompileConfig(), + ) + report = _operator_library_report(edge) + if not allow_cpu_fallback and report["cpu_ops"]: + raise RuntimeError( + "The ESP-NN operator library left these operators on the CPU: " + + ", ".join( + f"{op['method']}.{op['name']} ({op['target']})" + for op in report["cpu_ops"] + ) + ) + return edge, report + partitioner = EspNNPartitioner(allow_cpu_fallback=allow_cpu_fallback) edge = to_edge_transform_and_lower( program, diff --git a/backends/espressif/passes/__init__.py b/backends/espressif/passes/__init__.py new file mode 100644 index 00000000000..f5f66426af1 --- /dev/null +++ b/backends/espressif/passes/__init__.py @@ -0,0 +1,9 @@ +# 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 .esp_pass_manager import EspPassManager + +__all__ = ["EspPassManager"] diff --git a/backends/espressif/passes/esp_pass_manager.py b/backends/espressif/passes/esp_pass_manager.py new file mode 100644 index 00000000000..1c77850af81 --- /dev/null +++ b/backends/espressif/passes/esp_pass_manager.py @@ -0,0 +1,14 @@ +# 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.exir.pass_manager import ExportedProgramPassManager + + +class EspPassManager(ExportedProgramPassManager): + """Run Espressif operator-library transformations on an Edge program.""" + + def __init__(self) -> None: + super().__init__(passes=[]) diff --git a/backends/espressif/test/test_pass_manager.py b/backends/espressif/test/test_pass_manager.py new file mode 100644 index 00000000000..755cc9b9247 --- /dev/null +++ b/backends/espressif/test/test_pass_manager.py @@ -0,0 +1,185 @@ +# 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 import to_espnn_edge +from executorch.backends.espressif.passes import EspPassManager +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): + def forward(self, x, y): + return x * y + + +class MultiplyAndFloatUsers(torch.nn.Module): + def forward(self, x, y): + product = x * y + return product, torch.sin(x), torch.sin(product) + + +@pytest.mark.parametrize("entry_point", ["separate", "combined"]) +@pytest.mark.parametrize("broadcast", [False, True]) +def test_quantized_and_float_operations_survive_export(entry_point, broadcast): + 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) + ), + ) + 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) + edge = tester.get_artifact() + assert edge.compile_config._check_ir_validity + edge.exported_program().validate() + tester.to_executorch().serialize() + 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(): + inputs = (torch.linspace(-1, 1, 6), torch.linspace(-2, 2, 6)) + program = EspTester(Multiply(), inputs).quantize().export().get_artifact() + + delegated, delegate_report = to_espnn_edge(program, False) + assert delegate_report["partitions"] == 1 + assert delegate_report["delegated_ops"] == 1 + 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") + ] + assert edge.compile_config._check_ir_validity + torch.testing.assert_close( + edge.exported_program().module()(*inputs), + program.module()(*inputs), + atol=0, + rtol=0, + ) + serialized = deserialize_pte_binary(edge.to_executorch().buffer).program + assert not serialized.execution_plan[0].delegates + + +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\)"): + to_espnn_edge(program, False, use_operator_library=True) + + +def test_operator_library_strict_mode_accepts_no_compute(): + program = torch.export.export(torch.nn.Identity(), (torch.ones(3),)) + _, report = to_espnn_edge(program, False, use_operator_library=True) + assert report == {"esp_nn_ops": 0, "cpu_ops": []} + + +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() + ) + 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) + ) + serialized = deserialize_pte_binary(edge.to_executorch().buffer).program + assert not serialized.execution_plan[0].delegates + + +@pytest.mark.parametrize("entry_point", ["separate", "combined"]) +def test_passes_preserve_methods_and_constants(entry_point): + class Scale(torch.nn.Module): + def __init__(self, scale, offset): + super().__init__() + self.register_buffer("scale", scale) + self.offset = offset + + def forward(self, x): + return x * self.scale + self.offset + + inputs = (torch.tensor([-1.0, 2.0, 3.0]),) + models = { + "forward": Scale(torch.tensor([2.0, 3.0, 4.0]), torch.tensor([0.5, 1.0, 1.5])), + "other": Scale(torch.tensor([5.0, 6.0, 7.0]), torch.tensor([1.5, 2.0, 2.5])), + } + programs = { + name: torch.export.export(model, inputs) for name, model in models.items() + } + if entry_point == "combined": + transformed = to_edge_transform_and_lower( + programs, transform_passes=EspPassManager(), constant_methods={"version": 1} + ) + else: + stage = EspRunPasses() + stage.run(to_edge(programs, constant_methods={"version": 1})) + transformed = stage.artifact + assert transformed.methods == {"forward", "other"} + assert transformed.config_methods == {"version"} + for name, model in models.items(): + program = transformed.exported_program(name) + program.validate() + torch.testing.assert_close( + program.module()(*inputs), model(*inputs), atol=0, rtol=0 + ) + serialized = deserialize_pte_binary(transformed.to_executorch().buffer).program + assert {plan.name for plan in serialized.execution_plan} == { + "forward", + "other", + "version", + } + assert all(not plan.delegates for plan in serialized.execution_plan) + + +def test_operator_library_reports_all_methods(): + inputs = (torch.linspace(-1, 1, 6), torch.linspace(-2, 2, 6)) + programs = { + name: EspTester(Multiply(), inputs).quantize().export().get_artifact() + 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") + ] + for name, program in programs.items(): + torch.testing.assert_close( + edge.exported_program(name).module()(*inputs), + program.module()(*inputs), + atol=0, + rtol=0, + ) + serialized = deserialize_pte_binary(edge.to_executorch().buffer).program + assert {plan.name for plan in serialized.execution_plan} == {"forward", "other"} + assert all(not plan.delegates for plan in serialized.execution_plan) diff --git a/backends/espressif/test/tester.py b/backends/espressif/test/tester.py new file mode 100644 index 00000000000..c45cc300672 --- /dev/null +++ b/backends/espressif/test/tester.py @@ -0,0 +1,52 @@ +# 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.espressif.passes import EspPassManager +from executorch.backends.espressif.quantizer import EspNNQuantizer +from executorch.backends.test.harness import Tester +from executorch.backends.test.harness.stages import ( + Quantize, + RunPasses, + StageType, + ToEdgeTransformAndLower, +) +from executorch.exir import ( + EdgeCompileConfig, + EdgeProgramManager, + to_edge_transform_and_lower, +) + + +class EspRunPasses(RunPasses): + def __init__(self): + super().__init__(EspPassManager) + + def run(self, artifact: EdgeProgramManager, inputs=None) -> None: + self.edge_or_aten_program = artifact.transform(EspPassManager()) + + +class EspToEdgeTransformAndLower(ToEdgeTransformAndLower): + def __init__(self): + super().__init__(edge_compile_config=EdgeCompileConfig()) + + def run(self, artifact, inputs=None, generate_etrecord: bool = False) -> None: + self.edge_dialect_program = to_edge_transform_and_lower( + artifact, + compile_config=self.edge_compile_conf, + transform_passes=EspPassManager(), + generate_etrecord=generate_etrecord, + ) + + +class EspTester(Tester): + def __init__(self, module, example_inputs, calibration_samples=None): + stages = self.default_stage_classes() + stages[StageType.QUANTIZE] = lambda: Quantize( + EspNNQuantizer(), calibration_samples=calibration_samples + ) + stages[StageType.RUN_PASSES] = EspRunPasses + stages[StageType.TO_EDGE_TRANSFORM_AND_LOWER] = EspToEdgeTransformAndLower + super().__init__(module, example_inputs, stages)