Add opt-in Espressif operator-library export pipeline

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,

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"]

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=[])

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)

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)