Add Espressif multiply kernel and ESP-IDF tests

backends/espressif/CMakeLists.txt

@@ -0,0 +1,35 @@
+# 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.
+
+idf_component_register(
+  SRCS "ops/op_quantized_mul.cpp" PRIV_REQUIRES esp-nn WHOLE_ARCHIVE
+)
+
+get_filename_component(
+  EXECUTORCH_ROOT "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE
+)
+idf_build_get_property(PYTHON_EXECUTABLE PYTHON)
+include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake)
+
+set(_yaml_file "${CMAKE_CURRENT_SOURCE_DIR}/ops/operators.yaml")
+gen_selected_ops(LIB_NAME espressif_ops_lib OPS_SCHEMA_YAML "${_yaml_file}")
+generate_bindings_for_kernels(
+  LIB_NAME espressif_ops_lib CUSTOM_OPS_YAML "${_yaml_file}"
+)
+
+target_sources(
+  ${COMPONENT_LIB}
+  PRIVATE
+    "${CMAKE_CURRENT_BINARY_DIR}/espressif_ops_lib/RegisterCodegenUnboxedKernelsEverything.cpp"
+)
+target_include_directories(
+  ${COMPONENT_LIB} PRIVATE "${EXECUTORCH_ROOT}/.."
+                           "${EXECUTORCH_ROOT}/runtime/core/portable_type/c10"
+)
+target_compile_definitions(
+  ${COMPONENT_LIB} PRIVATE C10_USING_CUSTOM_GENERATED_MACROS
+)
+target_link_libraries(${COMPONENT_LIB} PUBLIC executorch_core)

backends/espressif/idf_component.yml

@@ -0,0 +1,11 @@
+# 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.
+
+description: ExecuTorch Espressif operator library
+dependencies:
+  idf:
+    version: ">=5.1"
+  espressif/esp-nn: "1.4.1"

backends/espressif/ops/op_quantized_mul.cpp

@@ -0,0 +1,63 @@
+/*
+ * 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.
+ */
+
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+#include <sdkconfig.h>
+
+#include <esp_nn.h>
+
+#include <cstdint>
+#include <limits>
+
+namespace espressif::native {
+
+using executorch::aten::Tensor;
+using executorch::runtime::KernelRuntimeContext;
+
+// cppcheck-suppress unusedFunction
+Tensor& quantized_mul_int8_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& input1,
+    const Tensor& input2,
+    int64_t input1_offset,
+    int64_t input2_offset,
+    int64_t output_offset,
+    int64_t output_multiplier,
+    int64_t output_shift,
+    Tensor& out) {
+  const auto numel = input1.numel();
+  ET_KERNEL_CHECK(
+      ctx,
+      numel >= 0 && numel <= std::numeric_limits<int32_t>::max(),
+      InvalidArgument,
+      out);
+  const auto size = static_cast<size_t>(numel);
+  ET_KERNEL_CHECK(
+      ctx,
+      size <= input1.nbytes() && size <= input2.nbytes() &&
+          size <= out.nbytes(),
+      InvalidArgument,
+      out);
+
+  esp_nn_mul_elementwise_s8(
+      input1.const_data_ptr<int8_t>(),
+      input2.const_data_ptr<int8_t>(),
+      static_cast<int32_t>(input1_offset),
+      static_cast<int32_t>(input2_offset),
+      out.mutable_data_ptr<int8_t>(),
+      static_cast<int32_t>(output_offset),
+      static_cast<int32_t>(output_multiplier),
+      static_cast<int32_t>(output_shift),
+      std::numeric_limits<int8_t>::min(),
+      std::numeric_limits<int8_t>::max(),
+      static_cast<int32_t>(numel));
+  return out;
+}
+
+} // namespace espressif::native

backends/espressif/ops/operators.yaml

@@ -0,0 +1,11 @@
+# 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.
+
+- func: esp_nn::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!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: espressif::quantized_mul_int8_out

backends/espressif/test/ops/test_mul.py

@@ -99,19 +99,44 @@ def forward(self, x, y):
         EspBroadcastMul(),
         (ramp_tensor(-1, 1, (2, 3)), ramp_tensor(-2, 2, (3,))),
     ),
+    "portable_add": EspTestCase(
+        EspMulAdd(),
+        (ramp_tensor(-1, 1, (17,)), ramp_tensor(2, -2, (17,))),
+    ),
+    "saturation": EspTestCase(
+        EspMul(),
+        (torch.tensor([-2.0, -2.0, 2.0, 2.0]), torch.tensor([-2.0, 2.0, -2.0, 2.0])),
+        calibration_samples=[
+            (
+                torch.tensor([-2.0, 2.0, 0.0, 0.0]),
+                torch.tensor([0.25, 0.25, -2.0, 2.0]),
+            )
+        ],
+    ),
 }
 
 
 @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 = EspTester(
+        test_case.model,
+        test_case.example_inputs,
+        calibration_samples=test_case.calibration_samples,
+    )
     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,
     )
+    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 node.args[:2]
+            )
     tester.to_executorch().serialize()
     program = deserialize_pte_binary(tester.get_artifact()).program
     assert not program.execution_plan[0].delegates
@@ -121,20 +146,13 @@ def test_dialect_mul(test_case):
     )
 
 
-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
+@pytest.mark.parametrize("test_case", test_cases.values(), ids=test_cases.keys())
+def test_implementation_mul(test_case, pytestconfig):
+    tester = EspTester(
+        test_case.model,
+        test_case.example_inputs,
+        calibration_samples=test_case.calibration_samples,
     )
-    assert all(
-        arg.target == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
-        for arg in add.args[:2]
+    tester.test_implementation(
+        pytestconfig.rootpath, rtol=torch.finfo(torch.float32).eps
     )
-    tester.to_executorch().serialize()

backends/espressif/test/tester.py

@@ -5,8 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 import math
+import os
+import subprocess
+import sys
 from dataclasses import dataclass
+from pathlib import Path
 
+import pytest
 import torch
 from executorch.backends.espressif.passes import EspPassManager
 from executorch.backends.espressif.quantizer import EspNNQuantizer
@@ -17,11 +22,17 @@
     StageType,
     ToEdgeTransformAndLower,
 )
+from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite
+from executorch.devtools.bundled_program.core import BundledProgram
+from executorch.devtools.bundled_program.serialize import (
+    serialize_from_bundled_program_to_flatbuffer,
+)
 from executorch.exir import (
     EdgeCompileConfig,
     EdgeProgramManager,
     to_edge_transform_and_lower,
 )
+from torch.utils._pytree import tree_flatten
 
 
 class EspRunPasses(RunPasses):
@@ -76,11 +87,94 @@ def test_dialect(
             )
         )
 
+    def test_implementation(self, repo: Path, *, atol=0, rtol=0):
+        build_path = os.environ.get("ESPRESSIF_BUILD_DIR")
+        if not build_path:
+            pytest.skip("Set ESPRESSIF_BUILD_DIR to the ESP32-S3 build directory")
+
+        self.quantize().export().to_edge().run_passes()
+        expected, _ = tree_flatten(
+            self.get_artifact().exported_program().module()(*self.example_inputs)
+        )
+        self.to_executorch().serialize()
+        bundle = BundledProgram(
+            self.get_artifact(StageType.TO_EXECUTORCH),
+            [
+                MethodTestSuite(
+                    "forward", [MethodTestCase(self.example_inputs, expected)]
+                )
+            ],
+        )
+
+        build_dir = Path(build_path).resolve()
+        worker = os.environ.get("PYTEST_XDIST_WORKER", "main")
+        firmware_dir = build_dir / f"operator-library-{worker}"
+        firmware_dir.mkdir(parents=True, exist_ok=True)
+        model_path = firmware_dir / "test.bpte"
+        model_path.write_bytes(serialize_from_bundled_program_to_flatbuffer(bundle))
+        commands = [
+            (
+                "build.log",
+                600,
+                [
+                    "idf.py",
+                    f"--project-dir={repo / 'examples/espressif/project'}",
+                    f"--build-dir={firmware_dir}",
+                    "-DIDF_TARGET=esp32s3",
+                    f"-DSDKCONFIG={firmware_dir / 'sdkconfig'}",
+                    "-DET_ESPRESSIF_OPS=ON",
+                    "-DET_ESPNN_BACKEND=OFF",
+                    f"-DET_BUILD_DIR_PATH={build_dir / 'runtime'}",
+                    f"-DET_PTE_FILE_PATH={model_path}",
+                    "-DET_BUNDLE_IO=ON",
+                    "-DET_NUM_INFERENCES=1",
+                    f"-DET_ATOL={atol}",
+                    f"-DET_RTOL={rtol}",
+                    "build",
+                    "merge-bin",
+                ],
+            ),
+            (
+                "emulator.log",
+                90,
+                [
+                    "esp-emu",
+                    "--chip=esp32s3",
+                    f"--firmware={firmware_dir / 'merged-binary.bin'}",
+                    f"--elf={firmware_dir / 'executorch_esp_runner.elf'}",
+                    "--net=user",
+                    "--psram-size=8M",
+                    "--timeout=60s",
+                    "--exit-on=Program complete.",
+                ],
+            ),
+            (
+                "check.log",
+                60,
+                [
+                    sys.executable,
+                    str(repo / "examples/espressif/test_emulator.py"),
+                    str(firmware_dir / "emulator.log"),
+                ],
+            ),
+        ]
+        for log_name, timeout, command in commands:
+            result = subprocess.run(
+                command,
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+            (firmware_dir / log_name).write_text(result.stdout + result.stderr)
+            assert result.returncode == 0, result.stdout + result.stderr
+        return self
+
 
 @dataclass
 class EspTestCase:
     model: torch.nn.Module
     example_inputs: tuple[torch.Tensor, ...]
+    calibration_samples: list[tuple[torch.Tensor, ...]] | None = None
 
 
 def ramp_tensor(start: float, end: float, shape: tuple[int, ...]) -> torch.Tensor:

examples/espressif/project/CMakeLists.txt

@@ -31,5 +31,10 @@ if(ET_ESPNN_BACKEND)
   )
 endif()
 
+option(ET_ESPRESSIF_OPS "Build the Espressif operator library" OFF)
+if(ET_ESPRESSIF_OPS)
+  list(APPEND EXTRA_COMPONENT_DIRS "${EXECUTORCH_ROOT}/backends/espressif")
+endif()
+
 include($ENV{IDF_PATH}/tools/cmake/project.cmake)
 project(executorch_esp_runner)