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)