From 000859ea9f7d142bc638d9ba6ca8e79076711c74 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 16 Oct 2026 10:00:45 -0400 Subject: [PATCH] Espressif: Add ESP-IDF component template for ExecuTorch Add the hand-written glue the component generator (next PR) stamps out: a registry-shaped CMakeLists.txt gated by CONFIG_ET_ENABLE with WHOLE_ARCHIVE kernel registration and size options (-Os, program verification, debug asserts, kernel registry size), the idf_component.yml and README templates, and et_model, a small C inference API over Program/Method whose method, planned and optional scratch arenas are caller-provided or placed with heap_caps. The component carries only the core runtime and BufferDataLoader, and the shared PAL is copied in by the generator rather than forked here. Authored with AI assistance using Claude Code. --- .../idf_component_template/CMakeLists.txt | 136 +++++ .../idf_component_template/README.md | 79 +++ .../idf_component.yml.in | 14 + .../idf_component_template/include/et_model.h | 179 ++++++ .../idf_component_template/src/et_model.cpp | 560 ++++++++++++++++++ 5 files changed, 968 insertions(+) create mode 100644 examples/espressif/idf_component_template/CMakeLists.txt create mode 100644 examples/espressif/idf_component_template/README.md create mode 100644 examples/espressif/idf_component_template/idf_component.yml.in create mode 100644 examples/espressif/idf_component_template/include/et_model.h create mode 100644 examples/espressif/idf_component_template/src/et_model.cpp diff --git a/examples/espressif/idf_component_template/CMakeLists.txt b/examples/espressif/idf_component_template/CMakeLists.txt new file mode 100644 index 00000000000..c5f346bc4d9 --- /dev/null +++ b/examples/espressif/idf_component_template/CMakeLists.txt @@ -0,0 +1,136 @@ +# 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. + +# ExecuTorch runtime as an ESP-IDF component. +# +# GENERATED by examples/espressif/build_idf_component.sh in the ExecuTorch repo +# -- do not edit this tree by hand. Source lists live in et_generated.cmake next +# to this file; the op profile is selected with CONFIG_ET_OP_PROFILE_* in +# Kconfig (no regeneration needed to switch). + +include(${CMAKE_CURRENT_LIST_DIR}/et_generated.cmake) + +set(_et_requires esp_hw_support esp_system esp_timer log heap) + +if(NOT CONFIG_ET_ENABLE) + # ESP-IDF links every component under components/ into the image, and this + # component's WHOLE_ARCHIVE registration (below) is deliberately GC-proof. The + # gate is off by default, so vendoring the component leaves a project's + # firmware unchanged until CONFIG_ET_ENABLE is set or selected. ESP-IDF's + # component-requirements scan runs this file before Kconfig is loaded, so it + # only ever sees this branch; registering the same requirements here keeps + # those components in the build when the component is enabled. No + # INCLUDE_DIRS: including while disabled fails at compile time + # rather than at link time. + idf_component_register(PRIV_REQUIRES ${_et_requires}) + return() +endif() + +# Pick the Kconfig-selected profile (default: ET_DEFAULT_PROFILE). +set(_et_profile "${ET_DEFAULT_PROFILE}") +foreach(_p IN LISTS ET_PROFILES) + string(TOUPPER "${_p}" _p_up) + if(CONFIG_ET_OP_PROFILE_${_p_up}) + set(_et_profile "${_p}") + endif() +endforeach() +string(TOUPPER "${_et_profile}" _ET_UP) + +set(_et_vendored_srcs ${ET_CORE_SRCS} ${ET_PROFILE_SRCS_${_ET_UP}} + ${ET_KERNEL_SRCS_${_ET_UP}} +) + +# WHOLE_ARCHIVE: kernel registration happens in static initializers in the +# profile's register_kernels.cpp; normal archive extraction would drop that +# translation unit (nothing references its symbols) and every operator would be +# missing at load time. --gc-sections still strips unused runtime and kernel +# code. src/platform/esp_pal.cpp is copied in by the generator from the shared +# examples/espressif/platform/esp_pal.cpp. +idf_component_register( + SRCS + ${_et_vendored_srcs} + "src/et_model.cpp" + "src/platform/esp_pal.cpp" + INCLUDE_DIRS + "include" + PRIV_REQUIRES + ${_et_requires} + WHOLE_ARCHIVE +) + +target_compile_features(${COMPONENT_LIB} PUBLIC cxx_std_17) + +set(_et_max_kernel_num ${ET_MAX_KERNEL_NUM_${_ET_UP}}) +if(CONFIG_ET_MAX_KERNEL_NUM) + if(CONFIG_ET_MAX_KERNEL_NUM LESS _et_max_kernel_num) + message( + FATAL_ERROR + "CONFIG_ET_MAX_KERNEL_NUM=${CONFIG_ET_MAX_KERNEL_NUM} is below the " + "${_et_profile} profile's ${_et_max_kernel_num} kernels" + ) + endif() + set(_et_max_kernel_num ${CONFIG_ET_MAX_KERNEL_NUM}) +endif() + +if(CONFIG_ET_PROGRAM_VERIFICATION) + set(_et_program_verification 1) +else() + set(_et_program_verification 0) +endif() + +# Consumers including headers need the same c10 macro mode and +# flatbuffers alignment the runtime was built with. __FUNCTION__ instead of +# __builtin_FUNCTION() keeps log function names without full signatures. +target_compile_definitions( + ${COMPONENT_LIB} + PUBLIC C10_USING_CUSTOM_GENERATED_MACROS FLATBUFFERS_MAX_ALIGNMENT=1024 + ET_USE_BUILTIN_FUNCTION_NAME=0 + PRIVATE MAX_KERNEL_NUM=${_et_max_kernel_num} + ET_ENABLE_PROGRAM_VERIFICATION=${_et_program_verification} +) + +if(NOT CONFIG_ET_DEBUG_ASSERTS) + target_compile_definitions(${COMPONENT_LIB} PRIVATE NDEBUG) +endif() + +if(CONFIG_ET_OPTIMIZE_SIZE) + target_compile_options(${COMPONENT_LIB} PRIVATE -Os) +endif() + +if(NOT CONFIG_ET_ENABLE_LOGGING) + target_compile_definitions(${COMPONENT_LIB} PUBLIC ET_LOG_ENABLED=0) +endif() + +# ET_MIN_LOG_LEVEL compiles out ET_LOG calls below the floor; ET_ESP_LOG_LEVEL +# applies the same floor in esp_pal.cpp. ExecuTorch has no warning level, so +# WARN keeps only errors. +if(CONFIG_ET_LOG_LEVEL_ERROR) + set(_et_min_log_level Error) + set(_et_esp_log_level ESP_LOG_ERROR) +elseif(CONFIG_ET_LOG_LEVEL_WARN) + set(_et_min_log_level Error) + set(_et_esp_log_level ESP_LOG_WARN) +elseif(CONFIG_ET_LOG_LEVEL_DEBUG) + set(_et_min_log_level Debug) + set(_et_esp_log_level ESP_LOG_DEBUG) +else() # CONFIG_ET_LOG_LEVEL_INFO (default) + set(_et_min_log_level Info) + set(_et_esp_log_level ESP_LOG_INFO) +endif() +target_compile_definitions( + ${COMPONENT_LIB} + PUBLIC ET_MIN_LOG_LEVEL=${_et_min_log_level} + PRIVATE ET_ESP_LOG_LEVEL=${_et_esp_log_level} +) + +# ESP-IDF builds with warnings-as-errors; upstream ExecuTorch sources are not +# warning-clean under the IDF flag set, so demote for vendored sources only. +set_source_files_properties( + ${_et_vendored_srcs} + PROPERTIES + COMPILE_OPTIONS + "-Wno-error;-Wno-deprecated-declarations;-Wno-unused-parameter;-Wno-sign-compare;-Wno-unused-variable" +) diff --git a/examples/espressif/idf_component_template/README.md b/examples/espressif/idf_component_template/README.md new file mode 100644 index 00000000000..09cee1aee59 --- /dev/null +++ b/examples/espressif/idf_component_template/README.md @@ -0,0 +1,79 @@ + + +# ExecuTorch for ESP-IDF (`executorch` component) + +**This tree is generated. Do not edit it by hand.** It is produced by +`examples/espressif/build_idf_component.sh` in the +[ExecuTorch](https://github.com/pytorch/executorch) repository: + +```sh +executorch/examples/espressif/build_idf_component.sh \ + --profiles @PROFILES@ --out +``` + +`VERSION` in this directory records the ExecuTorch commit, the op +profiles baked in, and each profile's op-set hash. A `.pte` is only +valid against the profile whose hash its export was checked against; +mismatches surface as `OperatorMissing` at load time. + +## Using it + +Add the component to your project's component requirements (`REQUIRES +executorch`), enable it (`CONFIG_ET_ENABLE=y`, or `select ET_ENABLE` +from your own Kconfig), pick an op profile in `menuconfig` (ExecuTorch -> Op +profile, or `CONFIG_ET_OP_PROFILE_=y` in your sdkconfig), then: + +- `#include ` for the small C inference API (flash XIP + model bytes; method, planned and optional scratch arenas either + passed in or allocated with heap_caps internal/PSRAM placement). This + is what most applications should use. The `.pte` bytes must be + 16-byte aligned; `et_model_create` rejects misaligned data. Application pipelines built on it (audio + front ends, detectors, wake word logic) live with the application, + not in this runtime component. +- `#include ` for the core C++ runtime + (`Program`/`Method`) and `BufferDataLoader`. `extension::Module`, the + file and mmap data loaders, and external `.ptd` data maps are not + included. Runtime logging goes through ESP_LOG under the + `executorch` tag. + +`CONFIG_ET_ENABLE` (default `n`) gates the whole component, so vendoring +it leaves a project's firmware unchanged until it is enabled. While +disabled the component registers empty and `et_model.h` is not on the +include path. + +## Op profiles + +@PROFILE_TABLE@ + +Only the selected profile's kernels are compiled in. Switching +profiles is a Kconfig change and a rebuild; no regeneration. + +## Size and debug options + +`CONFIG_ET_OPTIMIZE_SIZE` builds the runtime with `-Os` regardless of +the project's optimization level. `CONFIG_ET_PROGRAM_VERIFICATION` +compiles in full `.pte` verification, which `et_model_config_t.verify` +requires. `CONFIG_ET_DEBUG_ASSERTS` keeps ExecuTorch's debug +assertions (otherwise the runtime builds with `NDEBUG`). +`CONFIG_ET_MAX_KERNEL_NUM` raises the kernel registry size above the +profile's op count, for applications that register their own kernels. + +## Logging + +`CONFIG_ET_ENABLE_LOGGING=n` compiles ExecuTorch logging out entirely. +Otherwise runtime messages go to ESP_LOG (tag `executorch`), and +`CONFIG_ET_LOG_LEVEL_*` is a compile-time floor (`ET_MIN_LOG_LEVEL`) +on top of the usual IDF runtime levels. `et_model` logs its own +errors under the `et_model` tag. + +## Updating + +Re-run the generator against a newer ExecuTorch checkout and replace +this whole directory. Review the `VERSION` diff (commit, op-set +hashes, file counts, size); do not line-review vendored sources. diff --git a/examples/espressif/idf_component_template/idf_component.yml.in b/examples/espressif/idf_component_template/idf_component.yml.in new file mode 100644 index 00000000000..ab0f95d2935 --- /dev/null +++ b/examples/espressif/idf_component_template/idf_component.yml.in @@ -0,0 +1,14 @@ +## ExecuTorch runtime for ESP-IDF. +## +## GENERATED by examples/espressif/build_idf_component.sh -- do not edit. +## Version is "~"; see VERSION +## in this directory for the exact commit and op-set hashes. Vendored +## third-party code (FlatBuffers, flatcc: Apache-2.0) is listed in +## THIRD_PARTY_LICENSES. No targets list: the runtime is portable +## across ESP targets. +version: "@ET_COMPONENT_VERSION@" +description: ExecuTorch runtime for ESP-IDF +url: https://github.com/pytorch/executorch +license: BSD-3-Clause +dependencies: + idf: ">=6.0" diff --git a/examples/espressif/idf_component_template/include/et_model.h b/examples/espressif/idf_component_template/include/et_model.h new file mode 100644 index 00000000000..4dcd2ddd2ab --- /dev/null +++ b/examples/espressif/idf_component_template/include/et_model.h @@ -0,0 +1,179 @@ +/* + * 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. + */ + +/* + * et_model: a small C wrapper over the ExecuTorch runtime (Program/Method) + * for ESP-IDF applications that just want to run a .pte from flash. + * + * This header carries no ESP-IDF or C++ types. The model bytes are used + * in place (typically flash XIP) and must be 16-byte aligned. Method + * arenas are either caller-provided or allocated with heap_caps in + * et_model_load_method, and the inference path (set_input/execute/output) + * does not allocate unless a kernel requests scratch memory and there is + * no temp arena. + * + * et_model_config_t cfg = ET_MODEL_CONFIG_DEFAULT(model_pte, + * model_pte_size); + * cfg.planned_location = ET_MEM_INTERNAL; // always-on arenas in SRAM + * et_model_t *m = et_model_create(&cfg); + * if (!m || !et_model_load_method(m, NULL)) { handle error } + * et_model_set_input(m, 0, frame, sizeof(frame)); + * et_model_execute(m); + * size_t n; const void *out = et_model_output_data(m, 0, &n); + * + * Method state (mutable buffers, e.g. streaming-model history) persists + * across et_model_execute calls on the same loaded method, as in the C++ + * API. Loading a method again resets that state. + * + * Not thread-safe: use one et_model_t from one task, like the underlying + * executorch::runtime::Method. + */ + +#pragma once + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* Where an arena is allocated (mapped to heap_caps capabilities). */ +typedef enum { + ET_MEM_DEFAULT = 0, /* MALLOC_CAP_DEFAULT: IDF heap policy decides */ + ET_MEM_INTERNAL = 1, /* internal SRAM: fastest, scarcest */ + ET_MEM_PSRAM = 2, /* external PSRAM: fails if the board has none */ +} et_mem_location_t; + +typedef struct { + /* .pte bytes, 16-byte aligned; must outlive the et_model_t. */ + const uint8_t* pte; + size_t pte_size; + /* Each region below can instead use a caller-provided buffer: 16-byte + * aligned, outliving the et_model_t, and never freed by et_model. Its + * location is then ignored. */ + + /* Method allocator pool in bytes (method bookkeeping and non-planned + * tensors); 0 selects the default (64 KiB), and must be set when + * arena_buffer is. */ + size_t arena_size; + et_mem_location_t arena_location; + void* arena_buffer; + /* Scratch pool for kernels and delegates that request temporary memory + * (e.g. FFT, topk), reset after every kernel. 0 (default) means no + * pool; the runtime then mallocs scratch through the PAL during + * execute, only for kernels that need it. */ + size_t temp_arena_size; + et_mem_location_t temp_location; + void* temp_buffer; + /* Memory-planned activation/state buffers, zeroed on every method load + * so mutable state without exported initial values starts at zero. + * Always-on models usually want ET_MEM_INTERNAL; large models may need + * ET_MEM_PSRAM. A caller-provided planned_buffer must hold the total + * reported by et_model_planned_info. */ + et_mem_location_t planned_location; + void* planned_buffer; + size_t planned_buffer_size; + /* Fully verify the program flatbuffer at load. Costs load time; use it + * for .pte files that are not trusted at build time (e.g. OTA). + * Requires CONFIG_ET_PROGRAM_VERIFICATION; et_model_create fails + * otherwise. */ + bool verify; +} et_model_config_t; + +#define ET_MODEL_CONFIG_DEFAULT(pte_ptr, pte_len) \ + { \ + .pte = (pte_ptr), .pte_size = (pte_len), .arena_size = 0, \ + .arena_location = ET_MEM_DEFAULT, .arena_buffer = NULL, \ + .temp_arena_size = 0, .temp_location = ET_MEM_DEFAULT, \ + .temp_buffer = NULL, .planned_location = ET_MEM_DEFAULT, \ + .planned_buffer = NULL, .planned_buffer_size = 0, .verify = false, \ + } + +typedef struct et_model et_model_t; + +/* Loads the program (header check only unless cfg->verify is set). + * Returns NULL on failure (reason is logged via ESP_LOG). Does not load + * a method yet. */ +et_model_t* et_model_create(const et_model_config_t* cfg); +void et_model_destroy(et_model_t* model); + +size_t et_model_num_methods(et_model_t* model); +/* Name of method `index`, owned by the program; NULL if out of range. */ +const char* et_model_method_name(et_model_t* model, size_t index); + +/* Loads `name` (NULL = first method), allocating any arenas and planned + * buffers not provided by the caller. Replaces any previously loaded method and + * resets mutable state; on failure no method is loaded. Methods with more than + * 16 memory-planned buffers are rejected. */ +bool et_model_load_method(et_model_t* model, const char* name); + +/* Planned-buffer footprint of the named (NULL = loaded) method, for + * sizing memory budgets or planned_buffer without loading it. The total + * includes the padding that keeps each buffer 16-byte aligned. Returns false if + * the method does not exist, or if `name` is NULL and no method is + * loaded. Either out pointer may be NULL. */ +bool et_model_planned_info( + et_model_t* model, + const char* name, + size_t* out_num_buffers, + size_t* out_total_bytes); + +size_t et_model_num_inputs(et_model_t* model); +size_t et_model_num_outputs(et_model_t* model); + +/* Tensor metadata, mirroring ExecuTorch's TensorInfo. scalar_type uses + * the executorch::aten::ScalarType numbering (see + * executorch/runtime/core/portable_type/scalar_type.h). */ +#define ET_MODEL_MAX_DIMS 8 +typedef struct { + int32_t scalar_type; + size_t ndim; + int32_t dims[ET_MODEL_MAX_DIMS]; + size_t nbytes; +} et_tensor_info_t; + +/* Fail for non-tensor values and tensors with more than + * ET_MODEL_MAX_DIMS dims. */ +bool et_model_input_info( + et_model_t* model, + size_t index, + et_tensor_info_t* out); +bool et_model_output_info( + et_model_t* model, + size_t index, + et_tensor_info_t* out); + +/* Sets tensor input `index` from `data`. `nbytes` must equal the input's + * exported nbytes (et_model_input_info); dynamic shapes are not supported. + * Memory-planned inputs (the export default) are copied, so `data` may be + * reused immediately or live in flash. Inputs exported without memory + * planning keep pointing at `data`, which must then stay valid and + * unmodified until the input is set again. Non-tensor inputs cannot be + * set and keep their exported values. */ +bool et_model_set_input( + et_model_t* model, + size_t index, + const void* data, + size_t nbytes); + +bool et_model_execute(et_model_t* model); + +/* Output tensor storage after a successful execute; valid until the next + * execute/load/destroy. Returns NULL on error. *out_nbytes may be NULL. */ +const void* +et_model_output_data(et_model_t* model, size_t index, size_t* out_nbytes); + +/* Human-readable reason for the most recent failed call on `model`; + * never NULL. */ +const char* et_model_last_error(et_model_t* model); + +#ifdef __cplusplus +} +#endif diff --git a/examples/espressif/idf_component_template/src/et_model.cpp b/examples/espressif/idf_component_template/src/et_model.cpp new file mode 100644 index 00000000000..8cd55411561 --- /dev/null +++ b/examples/espressif/idf_component_template/src/et_model.cpp @@ -0,0 +1,560 @@ +/* + * 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. + */ + +// Implementation of et_model.h over executorch::runtime::Program/Method. +// +// Program/Method directly rather than extension::Module: Module keeps its +// methods in std::string-keyed maps and shared_ptr state on the global +// heap, while this API uses caller-provided or heap_caps buffers. Module +// is not vendored in this component. + +#include "et_model.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr char kTag[] = "et_model"; +constexpr size_t kDefaultArenaSize = 64 * 1024; +constexpr size_t kAlignment = 16; +constexpr size_t kMaxPlannedBuffers = 16; + +// program.cpp compiles the verifier in unless this is explicitly 0. +#if defined(ET_ENABLE_PROGRAM_VERIFICATION) && !ET_ENABLE_PROGRAM_VERIFICATION +constexpr bool kHasProgramVerification = false; +#else +constexpr bool kHasProgramVerification = true; +#endif + +using executorch::aten::Tensor; +using executorch::aten::TensorImpl; +using executorch::extension::BufferDataLoader; +using executorch::runtime::Error; +using executorch::runtime::EValue; +using executorch::runtime::HierarchicalAllocator; +using executorch::runtime::MemoryAllocator; +using executorch::runtime::MemoryManager; +using executorch::runtime::Method; +using executorch::runtime::MethodMeta; +using executorch::runtime::Program; +using executorch::runtime::Result; +using executorch::runtime::TensorInfo; +using SpanU8 = executorch::runtime::Span; + +bool valid_location(et_mem_location_t loc) { + return static_cast(loc) <= ET_MEM_PSRAM; +} + +bool is_aligned(const void* ptr) { + return reinterpret_cast(ptr) % kAlignment == 0; +} + +size_t align_up(size_t size) { + return (size + kAlignment - 1) / kAlignment * kAlignment; +} + +uint32_t caps_for(et_mem_location_t loc) { + switch (loc) { + case ET_MEM_INTERNAL: + return MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT; + case ET_MEM_PSRAM: + return MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT; + case ET_MEM_DEFAULT: + default: + return MALLOC_CAP_DEFAULT; + } +} + +// Zero-sized planned buffers still need a valid pointer for their span. +uint8_t* caps_alloc(size_t size, et_mem_location_t loc) { + return static_cast( + heap_caps_aligned_alloc(kAlignment, size == 0 ? 1 : size, caps_for(loc))); +} + +} // namespace + +struct et_model { + explicit et_model(const et_model_config_t& c) + : cfg(c), loader(c.pte, c.pte_size) {} + + void fail(const char* what, Error err) { + snprintf( + error, + sizeof(error), + "%s failed (error 0x%x)", + what, + static_cast(err)); + ESP_LOGE(kTag, "%s", error); + } + void fail(const char* what) { + snprintf(error, sizeof(error), "%s", what); + ESP_LOGE(kTag, "%s", error); + } + + void free_method() { + meta.reset(); + method.reset(); + manager.reset(); + planned.reset(); + for (size_t i = 0; i < num_planned_owned; i++) { + heap_caps_free(planned_ptrs[i]); + planned_ptrs[i] = nullptr; + } + num_planned_owned = 0; + temp_allocator.reset(); + heap_caps_free(temp_arena); + temp_arena = nullptr; + allocator.reset(); + heap_caps_free(arena); + arena = nullptr; + } + + et_model_config_t cfg; + BufferDataLoader loader; + std::optional program; + + // Loaded-method state (all zero/null until et_model_load_method). arena, + // temp_arena and planned_ptrs hold only allocations owned by et_model. + uint8_t* arena = nullptr; + std::optional allocator; + uint8_t* temp_arena = nullptr; + std::optional temp_allocator; + uint8_t* planned_ptrs[kMaxPlannedBuffers] = {}; + size_t num_planned_owned = 0; + SpanU8 planned_spans[kMaxPlannedBuffers]; + std::optional planned; + std::optional manager; + std::optional method; + std::optional meta; + + char error[128] = "no error"; +}; + +namespace { + +bool fill_info( + et_model* model, + const Result& info, + et_tensor_info_t* out) { + if (!info.ok()) { + model->fail("tensor_meta", info.error()); + return false; + } + if (out == nullptr) { + model->fail("null et_tensor_info_t"); + return false; + } + if (info->sizes().size() > ET_MODEL_MAX_DIMS) { + model->fail("tensor has more than ET_MODEL_MAX_DIMS dims"); + return false; + } + memset(out, 0, sizeof(*out)); + out->scalar_type = static_cast(info->scalar_type()); + out->ndim = info->sizes().size(); + for (size_t i = 0; i < out->ndim; i++) { + out->dims[i] = info->sizes()[i]; + } + out->nbytes = info->nbytes(); + return true; +} + +// Resolves `name` (NULL = first method). The returned name is either +// `name` or owned by the program. +bool method_meta_for( + et_model* model, + const char* name, + std::optional* out_meta, + const char** out_name) { + const char* resolved = name; + if (resolved == nullptr) { + auto name_res = model->program->get_method_name(0); + if (!name_res.ok()) { + model->fail("program has no methods"); + return false; + } + resolved = *name_res; + } + auto meta_res = model->program->method_meta(resolved); + if (!meta_res.ok()) { + model->fail("method_meta", meta_res.error()); + return false; + } + out_meta->emplace(std::move(meta_res.get())); + *out_name = resolved; + return true; +} + +bool require_method(et_model* model) { + if (!model->method) { + model->fail("no method loaded"); + return false; + } + return true; +} + +} // namespace + +extern "C" { + +et_model_t* et_model_create(const et_model_config_t* cfg) { + if (cfg == nullptr || cfg->pte == nullptr || cfg->pte_size == 0) { + ESP_LOGE(kTag, "et_model_create: invalid config"); + return nullptr; + } + if (!is_aligned(cfg->pte) || !is_aligned(cfg->arena_buffer) || + !is_aligned(cfg->temp_buffer) || !is_aligned(cfg->planned_buffer)) { + ESP_LOGE( + kTag, + "et_model_create: pte and caller buffers must be %zu-byte aligned", + kAlignment); + return nullptr; + } + if ((cfg->arena_buffer != nullptr && cfg->arena_size == 0) || + (cfg->temp_buffer != nullptr && cfg->temp_arena_size == 0) || + (cfg->planned_buffer != nullptr && cfg->planned_buffer_size == 0)) { + ESP_LOGE(kTag, "et_model_create: caller buffer given without its size"); + return nullptr; + } + if (!valid_location(cfg->arena_location) || + !valid_location(cfg->temp_location) || + !valid_location(cfg->planned_location)) { + ESP_LOGE(kTag, "et_model_create: invalid et_mem_location_t"); + return nullptr; + } + if (cfg->verify && !kHasProgramVerification) { + ESP_LOGE( + kTag, + "et_model_create: verify requires CONFIG_ET_PROGRAM_VERIFICATION"); + return nullptr; + } + executorch::runtime::runtime_init(); + + std::unique_ptr model(new (std::nothrow) et_model(*cfg)); + if (!model) { + ESP_LOGE(kTag, "et_model_create: out of memory"); + return nullptr; + } + auto program_res = Program::load( + &model->loader, + cfg->verify ? Program::Verification::InternalConsistency + : Program::Verification::Minimal); + if (!program_res.ok()) { + model->fail("Program::load", program_res.error()); + return nullptr; + } + model->program.emplace(std::move(program_res.get())); + ESP_LOGI( + kTag, + "loaded program: %zu bytes, %zu method(s)", + cfg->pte_size, + static_cast(model->program->num_methods())); + return model.release(); +} + +void et_model_destroy(et_model_t* model) { + if (model == nullptr) { + return; + } + model->free_method(); + delete model; +} + +size_t et_model_num_methods(et_model_t* model) { + if (model == nullptr) { + return 0; + } + return static_cast(model->program->num_methods()); +} + +const char* et_model_method_name(et_model_t* model, size_t index) { + if (model == nullptr) { + return nullptr; + } + auto res = model->program->get_method_name(index); + if (!res.ok()) { + model->fail("get_method_name", res.error()); + return nullptr; + } + return *res; +} + +bool et_model_load_method(et_model_t* model, const char* name) { + if (model == nullptr) { + return false; + } + model->free_method(); + + std::optional meta; + const char* resolved = nullptr; + if (!method_meta_for(model, name, &meta, &resolved)) { + return false; + } + const size_t num_planned = meta->num_memory_planned_buffers(); + if (num_planned > kMaxPlannedBuffers) { + model->fail("too many planned buffers"); + return false; + } + + const size_t arena_size = + model->cfg.arena_size ? model->cfg.arena_size : kDefaultArenaSize; + uint8_t* arena = static_cast(model->cfg.arena_buffer); + if (arena == nullptr) { + arena = model->arena = caps_alloc(arena_size, model->cfg.arena_location); + if (arena == nullptr) { + model->fail("allocating method arena failed"); + return false; + } + } + model->allocator.emplace(static_cast(arena_size), arena); + + const size_t temp_arena_size = model->cfg.temp_arena_size; + if (temp_arena_size != 0) { + uint8_t* temp = static_cast(model->cfg.temp_buffer); + if (temp == nullptr) { + temp = model->temp_arena = + caps_alloc(temp_arena_size, model->cfg.temp_location); + if (temp == nullptr) { + model->free_method(); + model->fail("allocating temp arena failed"); + return false; + } + } + model->temp_allocator.emplace(static_cast(temp_arena_size), temp); + } + + uint8_t* planned_base = static_cast(model->cfg.planned_buffer); + size_t planned_offset = 0; + for (size_t i = 0; i < num_planned; i++) { + auto size_res = meta->memory_planned_buffer_size(i); + if (!size_res.ok()) { + model->free_method(); + model->fail("memory_planned_buffer_size", size_res.error()); + return false; + } + const size_t bytes = static_cast(size_res.get()); + uint8_t* buf = nullptr; + if (planned_base != nullptr) { + if (align_up(bytes) > model->cfg.planned_buffer_size - planned_offset) { + model->free_method(); + model->fail("planned_buffer smaller than et_model_planned_info total"); + return false; + } + buf = planned_base + planned_offset; + planned_offset += align_up(bytes); + } else { + buf = caps_alloc(bytes, model->cfg.planned_location); + if (buf == nullptr) { + model->free_method(); + model->fail("allocating planned buffer failed"); + return false; + } + model->planned_ptrs[i] = buf; + model->num_planned_owned = i + 1; + } + // Mutable buffers exported without initial values start from whatever + // the planned memory holds. + memset(buf, 0, bytes); + model->planned_spans[i] = SpanU8(buf, bytes); + } + model->planned.emplace( + executorch::runtime::Span(model->planned_spans, num_planned)); + // A null temp allocator makes Method fall back to PAL allocation. + model->manager.emplace( + &*model->allocator, + &*model->planned, + model->temp_allocator ? &*model->temp_allocator : nullptr); + + auto method_res = model->program->load_method(resolved, &*model->manager); + if (!method_res.ok()) { + Error err = method_res.error(); + model->free_method(); + model->fail("load_method", err); + return false; + } + model->method.emplace(std::move(method_res.get())); + model->meta.emplace(std::move(*meta)); + ESP_LOGI( + kTag, + "loaded method '%s': arena %zu B, temp %zu B, %zu planned buffer(s)", + resolved, + arena_size, + temp_arena_size, + num_planned); + return true; +} + +bool et_model_planned_info( + et_model_t* model, + const char* name, + size_t* out_num_buffers, + size_t* out_total_bytes) { + if (model == nullptr) { + return false; + } + std::optional meta; + if (name == nullptr) { + if (!model->meta) { + model->fail("no method loaded"); + return false; + } + meta = model->meta; + } else { + const char* resolved = nullptr; + if (!method_meta_for(model, name, &meta, &resolved)) { + return false; + } + } + size_t total = 0; + for (size_t i = 0; i < meta->num_memory_planned_buffers(); i++) { + auto size_res = meta->memory_planned_buffer_size(i); + if (!size_res.ok()) { + model->fail("memory_planned_buffer_size", size_res.error()); + return false; + } + total += align_up(static_cast(size_res.get())); + } + if (out_num_buffers != nullptr) { + *out_num_buffers = meta->num_memory_planned_buffers(); + } + if (out_total_bytes != nullptr) { + *out_total_bytes = total; + } + return true; +} + +size_t et_model_num_inputs(et_model_t* model) { + return (model != nullptr && model->meta) ? model->meta->num_inputs() : 0; +} + +size_t et_model_num_outputs(et_model_t* model) { + return (model != nullptr && model->meta) ? model->meta->num_outputs() : 0; +} + +bool et_model_input_info( + et_model_t* model, + size_t index, + et_tensor_info_t* out) { + if (model == nullptr || !require_method(model)) { + return false; + } + return fill_info(model, model->meta->input_tensor_meta(index), out); +} + +bool et_model_output_info( + et_model_t* model, + size_t index, + et_tensor_info_t* out) { + if (model == nullptr || !require_method(model)) { + return false; + } + return fill_info(model, model->meta->output_tensor_meta(index), out); +} + +bool et_model_set_input( + et_model_t* model, + size_t index, + const void* data, + size_t nbytes) { + if (model == nullptr || !require_method(model)) { + return false; + } + auto info = model->meta->input_tensor_meta(index); + if (!info.ok()) { + model->fail("input_tensor_meta", info.error()); + return false; + } + if (info->nbytes() != nbytes) { + model->fail("input size does not match model"); + return false; + } + const size_t ndim = info->sizes().size(); + if (ndim > ET_MODEL_MAX_DIMS) { + model->fail("input has more than ET_MODEL_MAX_DIMS dims"); + return false; + } + // Method copies planned inputs and shares only the data pointer of + // non-planned ones, so the tensor wrapper can live on the stack. + int32_t sizes[ET_MODEL_MAX_DIMS]; + uint8_t dim_order[ET_MODEL_MAX_DIMS]; + for (size_t d = 0; d < ndim; d++) { + sizes[d] = info->sizes()[d]; + dim_order[d] = info->dim_order()[d]; + } + TensorImpl impl( + info->scalar_type(), + static_cast(ndim), + sizes, + const_cast(data), + dim_order); + Error err = model->method->set_input(EValue(Tensor(&impl)), index); + if (err != Error::Ok) { + model->fail("set_input", err); + return false; + } + return true; +} + +bool et_model_execute(et_model_t* model) { + if (model == nullptr || !require_method(model)) { + return false; + } + Error err = model->method->execute(); + if (err != Error::Ok) { + model->fail("execute", err); + return false; + } + return true; +} + +const void* +et_model_output_data(et_model_t* model, size_t index, size_t* out_nbytes) { + if (out_nbytes != nullptr) { + *out_nbytes = 0; + } + if (model == nullptr || !require_method(model)) { + return nullptr; + } + if (index >= model->method->outputs_size()) { + model->fail("output index out of range"); + return nullptr; + } + const EValue& output = model->method->get_output(index); + if (!output.isTensor()) { + model->fail("output is not a tensor"); + return nullptr; + } + const Tensor& tensor = output.toTensor(); + if (out_nbytes != nullptr) { + *out_nbytes = tensor.nbytes(); + } + return tensor.const_data_ptr(); +} + +const char* et_model_last_error(et_model_t* model) { + return model != nullptr ? model->error : "null model"; +} + +} // extern "C"