diff --git a/firmware/esp32_tinystories_idf/.gitignore b/firmware/esp32_tinystories_idf/.gitignore new file mode 100644 index 0000000..48fd694 --- /dev/null +++ b/firmware/esp32_tinystories_idf/.gitignore @@ -0,0 +1,7 @@ +/build/ +/managed_components/ +/components/gen_bmgr_codes/ +/generated/ +/sdkconfig +/sdkconfig.old +/dependencies.lock diff --git a/firmware/esp32_tinystories_idf/CMakeLists.txt b/firmware/esp32_tinystories_idf/CMakeLists.txt new file mode 100644 index 0000000..d376346 --- /dev/null +++ b/firmware/esp32_tinystories_idf/CMakeLists.txt @@ -0,0 +1,22 @@ +cmake_minimum_required(VERSION 3.16) + +# `idf.py bmgr` generates board-specific defaults here. Load them first so the +# project's target-specific performance defaults can override conservative +# board memory settings. +set(BMGR_DEFAULTS "${CMAKE_CURRENT_LIST_DIR}/components/gen_bmgr_codes/board_manager.defaults") +if(EXISTS "${BMGR_DEFAULTS}") + list(FIND SDKCONFIG_DEFAULTS "${BMGR_DEFAULTS}" BMGR_DEFAULTS_INDEX) + if(BMGR_DEFAULTS_INDEX EQUAL -1) + if(SDKCONFIG_DEFAULTS) + list(PREPEND SDKCONFIG_DEFAULTS "${BMGR_DEFAULTS}") + else() + set(SDKCONFIG_DEFAULTS + "${BMGR_DEFAULTS};${CMAKE_CURRENT_LIST_DIR}/sdkconfig.defaults" + ) + endif() + endif() +endif() + +include($ENV{IDF_PATH}/tools/cmake/project.cmake) + +project(esp32_tinystories_idf) diff --git a/firmware/esp32_tinystories_idf/README.md b/firmware/esp32_tinystories_idf/README.md new file mode 100644 index 0000000..a533b00 --- /dev/null +++ b/firmware/esp32_tinystories_idf/README.md @@ -0,0 +1,145 @@ +# TinyStories ESP-IDF on-chip inference + +This project is an ESP-IDF and ESP Board Manager port of [`firmware/esp32_tinystories`](../esp32_tinystories/README.md). + +It runs the released 28.9M-parameter TinyStories PLE model and shares the portable inference runtime in [`runtime/llm.h`](../../runtime/llm.h) with the upstream Arduino firmware and host verification tools. + +The application is not tied to one board because Board Manager supplies the target, memory configuration, display handle, display resolution, and optional backlight device. + +## Project structure + +- `main/esp32_tinystories_idf.cpp` contains model loading, memory placement, inference, profiling, and text generation. +- `main/board_display.cpp` renders generated text through the Board Manager `display_lcd` device. +- `generated/vocab.h` is generated from the tokenizer being deployed and is intentionally excluded from Git. +- `../../runtime/llm.h` is the shared PLE model loader and inference implementation. +- `../../artifacts/tinystories/model.bin` is the verified released model downloaded by the repository fetch script. + +## Runtime layout + +- Flash holds the memory-mapped token embedding and PLE table, which are accessed by row. +- PSRAM holds the staged int8 core and output-head weights, logits, and KV cache. +- Internal SRAM holds the hot scratch buffers and relocated RMSNorm vectors. +- FreeRTOS task notifications split sufficiently large staged matvec operations across both cores when the selected target has more than one core. +- Generated text is written to both the serial console and the LCD resolution reported by Board Manager. + +The implementation uses the current published TinyStories runtime with int8 activations and stages both the dense core and output head instead of retaining the earlier output-head-only optimization. + +## Requirements + +- ESP-IDF 5.5 or later. +- Python package `esp-bmgr-assist` for the `idf.py bmgr` action. +- Hugging Face CLI command `hf` for the verified model fetch script. +- A supported Board Manager definition with at least 16 MB Flash, PSRAM, and an RGB565 `display_lcd` device. + +## Prepare the released model + +Training is not required for normal deployment because the repository publishes a verified TinyStories model and tokenizer. + +Run the following command from the repository root to download, verify, and install the release under `artifacts/tinystories`: + +```bash +scripts/fetch_model.sh tinystories +``` + +The fetch script validates the model and tokenizer against pinned byte sizes and SHA-256 values before replacing any local artifact. + +Generate this ESP-IDF project's decode header from the tokenizer that will be deployed: + +```bash +uv run python firmware/esp32_tinystories/tools/generate_vocab.py \ + --tokenizer artifacts/tinystories/tokenizer.json \ + --out firmware/esp32_tinystories_idf/generated/vocab.h +``` + +The generated header remains inside `firmware/esp32_tinystories_idf` and is ignored by Git. + +Training and `research.tinystories.export` are only needed when reproducing or replacing the published model. + +## Select the board + +Load the ESP-IDF environment and install the Board Manager action helper once: + +```bash +. "$IDF_PATH/export.sh" +python -m pip install esp-bmgr-assist +``` + +Enter this project, list the available boards, and generate the selected board configuration: + +```bash +cd firmware/esp32_tinystories_idf +idf.py bmgr -l +idf.py bmgr -b +``` + +Select ESP32-P4-Function-EV-Board with: + +```bash +idf.py bmgr -b esp32_p4_function_ev_board +``` + +Select ESP32-S31-Korvo-1 with: + +```bash +idf.py bmgr -b esp32_s31_korvo_1 +``` + +Select ESP32-S3-LCD-EV-Board with its default 480x480 GC9503 LCD sub-board with: + +```bash +idf.py bmgr -b esp32_s3_lcd_ev_board +``` + +Select its 800x480 ST7262 RGB LCD and GT1151 touch sub-board, whose silkscreen is ESP32-S3-LCD-EV-Board-SUB3, with: + +```bash +idf.py bmgr -b esp32_s3_lcd_ev_board \ + -a sub_board_800_480_lcd +``` + +The `bmgr` action selects the board's SoC target and generates `components/gen_bmgr_codes`, so a separate `idf.py set-target` command is not required. + +## Build and flash + +Build the firmware with: + +```bash +idf.py build +``` + +Flash the bootloader, partition table, application, and released model, then open the serial monitor: + +```bash +idf.py -p PORT flash monitor +``` + +The full `flash` target always includes `artifacts/tinystories/model.bin`, so flashing fails before writing the board when the model has not been fetched. + +After the model has been installed once, firmware-only changes can be updated without rewriting the model partition: + +```bash +idf.py -p PORT app-flash monitor +``` + +Press `Ctrl-]` to exit the serial monitor. + +## Board Manager initialization + +At startup, the application performs the following operations: + +1. `esp_board_manager_print_board_info()` prints the selected board metadata. +2. `esp_board_manager_init_device_by_name("display_lcd")` initializes the selected board's panel and returns a generic `esp_lcd_panel_handle_t`. +3. The optional `lcd_brightness` device is initialized and set to full brightness when it is present. +4. The `model` partition is memory-mapped and checked against the generated tokenizer vocabulary. +5. Hot scratch and normalization data are placed in internal SRAM while staged weights, logits, and cache data are placed in PSRAM. +6. The application enables dual-core matvec execution when the selected target provides more than one core. + +## Regenerate the board configuration + +Run the following sequence when changing boards or when stale `sdkconfig` values prevent new board defaults from taking effect: + +```bash +idf.py bmgr -x +idf.py fullclean +idf.py bmgr -b +``` diff --git a/firmware/esp32_tinystories_idf/main/CMakeLists.txt b/firmware/esp32_tinystories_idf/main/CMakeLists.txt new file mode 100644 index 0000000..245ea0d --- /dev/null +++ b/firmware/esp32_tinystories_idf/main/CMakeLists.txt @@ -0,0 +1,43 @@ +set(TINYSTORIES_RUNTIME_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../runtime") +set(TINYSTORIES_VOCAB_DIR "${CMAKE_CURRENT_LIST_DIR}/../generated") +set(TINYSTORIES_VOCAB_HEADER "${TINYSTORIES_VOCAB_DIR}/vocab.h") +set( + TINYSTORIES_ARTIFACT_DIR + "${CMAKE_CURRENT_LIST_DIR}/../../../artifacts/tinystories" + CACHE PATH + "Directory containing the verified TinyStories release artifacts" +) +set(TINYSTORIES_MODEL_BIN "${TINYSTORIES_ARTIFACT_DIR}/model.bin") + +if(NOT EXISTS "${TINYSTORIES_RUNTIME_DIR}/llm.h") + message(FATAL_ERROR "Missing shared runtime: ${TINYSTORIES_RUNTIME_DIR}/llm.h") +endif() + +if(NOT EXISTS "${TINYSTORIES_VOCAB_HEADER}") + message(FATAL_ERROR "Missing generated vocabulary: ${TINYSTORIES_VOCAB_HEADER}. Follow the Prepare the released model section in firmware/esp32_tinystories_idf/README.md.") +endif() + +idf_component_register( + SRCS "esp32_tinystories_idf.cpp" "board_display.cpp" + INCLUDE_DIRS "." "${TINYSTORIES_RUNTIME_DIR}" "${TINYSTORIES_VOCAB_DIR}" + REQUIRES esp_board_manager esp_lcd esp_partition esp_timer driver heap +) + +target_compile_options(${COMPONENT_LIB} PRIVATE -O3) + +# A full `idf.py flash` always includes the released model. This intentionally +# makes flashing fail before touching the board when the artifact is absent. +# `app-flash` remains available after the model has been installed once. +set(TINYSTORIES_MODEL_PARTITION_SIZE 15007744) # 0xE50000 +if(EXISTS "${TINYSTORIES_MODEL_BIN}") + file(SIZE "${TINYSTORIES_MODEL_BIN}" TINYSTORIES_MODEL_BIN_SIZE) + if(TINYSTORIES_MODEL_BIN_SIZE GREATER TINYSTORIES_MODEL_PARTITION_SIZE) + message(FATAL_ERROR + "model.bin (${TINYSTORIES_MODEL_BIN_SIZE} bytes) exceeds the model partition " + "(${TINYSTORIES_MODEL_PARTITION_SIZE} bytes)" + ) + endif() +else() + message(WARNING "artifacts/tinystories/model.bin is absent. Run scripts/fetch_model.sh tinystories before flashing.") +endif() +esptool_py_flash_to_partition(flash "model" "${TINYSTORIES_MODEL_BIN}") diff --git a/firmware/esp32_tinystories_idf/main/board_display.cpp b/firmware/esp32_tinystories_idf/main/board_display.cpp new file mode 100644 index 0000000..32c3948 --- /dev/null +++ b/firmware/esp32_tinystories_idf/main/board_display.cpp @@ -0,0 +1,317 @@ +#include "board_display.h" + +#include +#include +#include + +#include "driver/ledc.h" +#include "esp_board_manager_includes.h" +#include "esp_check.h" +#include "esp_heap_caps.h" +#include "esp_lcd_panel_ops.h" +#include "esp_log.h" + +namespace { + +constexpr uint16_t kBackground = 0x0000; // RGB565 black +constexpr uint16_t kForeground = 0x07E0; // RGB565 green +constexpr int kGlyphWidth = 5; +constexpr int kGlyphHeight = 7; + +const char *const TAG = "llm_display"; + +esp_lcd_panel_handle_t panel; +uint16_t *framebuffer; +uint16_t width; +uint16_t height; +uint16_t cursor_x; +uint16_t cursor_y; +uint8_t scale; +bool display_ready; + +// Compact 5x7 font. Lowercase is deliberately rendered as uppercase: the +// TinyStories output stays readable while the complete font costs only 180 B. +constexpr uint8_t kDigits[10][kGlyphWidth] = { + {0x3E, 0x51, 0x49, 0x45, 0x3E}, {0x00, 0x42, 0x7F, 0x40, 0x00}, + {0x42, 0x61, 0x51, 0x49, 0x46}, {0x21, 0x41, 0x45, 0x4B, 0x31}, + {0x18, 0x14, 0x12, 0x7F, 0x10}, {0x27, 0x45, 0x45, 0x45, 0x39}, + {0x3C, 0x4A, 0x49, 0x49, 0x30}, {0x01, 0x71, 0x09, 0x05, 0x03}, + {0x36, 0x49, 0x49, 0x49, 0x36}, {0x06, 0x49, 0x49, 0x29, 0x1E}, +}; + +constexpr uint8_t kLetters[26][kGlyphWidth] = { + {0x7E, 0x11, 0x11, 0x11, 0x7E}, {0x7F, 0x49, 0x49, 0x49, 0x36}, + {0x3E, 0x41, 0x41, 0x41, 0x22}, {0x7F, 0x41, 0x41, 0x22, 0x1C}, + {0x7F, 0x49, 0x49, 0x49, 0x41}, {0x7F, 0x09, 0x09, 0x09, 0x01}, + {0x3E, 0x41, 0x49, 0x49, 0x7A}, {0x7F, 0x08, 0x08, 0x08, 0x7F}, + {0x00, 0x41, 0x7F, 0x41, 0x00}, {0x20, 0x40, 0x41, 0x3F, 0x01}, + {0x7F, 0x08, 0x14, 0x22, 0x41}, {0x7F, 0x40, 0x40, 0x40, 0x40}, + {0x7F, 0x02, 0x0C, 0x02, 0x7F}, {0x7F, 0x04, 0x08, 0x10, 0x7F}, + {0x3E, 0x41, 0x41, 0x41, 0x3E}, {0x7F, 0x09, 0x09, 0x09, 0x06}, + {0x3E, 0x41, 0x51, 0x21, 0x5E}, {0x7F, 0x09, 0x19, 0x29, 0x46}, + {0x46, 0x49, 0x49, 0x49, 0x31}, {0x01, 0x01, 0x7F, 0x01, 0x01}, + {0x3F, 0x40, 0x40, 0x40, 0x3F}, {0x1F, 0x20, 0x40, 0x20, 0x1F}, + {0x3F, 0x40, 0x38, 0x40, 0x3F}, {0x63, 0x14, 0x08, 0x14, 0x63}, + {0x07, 0x08, 0x70, 0x08, 0x07}, {0x61, 0x51, 0x49, 0x45, 0x43}, +}; + +void punctuation_glyph(char c, uint8_t out[kGlyphWidth]) { + const uint8_t *glyph = nullptr; + static constexpr uint8_t exclamation[] = {0x00, 0x00, 0x5F, 0x00, 0x00}; + static constexpr uint8_t quote[] = {0x00, 0x07, 0x00, 0x07, 0x00}; + static constexpr uint8_t apostrophe[] = {0x00, 0x05, 0x03, 0x00, 0x00}; + static constexpr uint8_t comma[] = {0x00, 0x50, 0x30, 0x00, 0x00}; + static constexpr uint8_t dash[] = {0x08, 0x08, 0x08, 0x08, 0x08}; + static constexpr uint8_t period[] = {0x00, 0x60, 0x60, 0x00, 0x00}; + static constexpr uint8_t colon[] = {0x00, 0x36, 0x36, 0x00, 0x00}; + static constexpr uint8_t semicolon[] = {0x00, 0x56, 0x36, 0x00, 0x00}; + static constexpr uint8_t question[] = {0x02, 0x01, 0x51, 0x09, 0x06}; + static constexpr uint8_t slash[] = {0x20, 0x10, 0x08, 0x04, 0x02}; + static constexpr uint8_t paren_l[] = {0x00, 0x1C, 0x22, 0x41, 0x00}; + static constexpr uint8_t paren_r[] = {0x00, 0x41, 0x22, 0x1C, 0x00}; + switch (c) { + case '!': + glyph = exclamation; + break; + case '"': + glyph = quote; + break; + case '\'': + glyph = apostrophe; + break; + case ',': + glyph = comma; + break; + case '-': + glyph = dash; + break; + case '.': + glyph = period; + break; + case ':': + glyph = colon; + break; + case ';': + glyph = semicolon; + break; + case '?': + glyph = question; + break; + case '/': + glyph = slash; + break; + case '(': + glyph = paren_l; + break; + case ')': + glyph = paren_r; + break; + default: + glyph = question; + break; + } + memcpy(out, glyph, kGlyphWidth); +} + +void get_glyph(char c, uint8_t out[kGlyphWidth]) { + if (c == ' ') { + memset(out, 0, kGlyphWidth); + } else if (c >= '0' && c <= '9') { + memcpy(out, kDigits[c - '0'], kGlyphWidth); + } else { + c = static_cast(toupper(static_cast(c))); + if (c >= 'A' && c <= 'Z') { + memcpy(out, kLetters[c - 'A'], kGlyphWidth); + } else { + punctuation_glyph(c, out); + } + } +} + +void clear_screen() { + memset(framebuffer, 0, + static_cast(width) * height * sizeof(uint16_t)); + cursor_x = scale * 2; + cursor_y = scale * 2; +} + +void newline() { + cursor_x = scale * 2; + cursor_y += (kGlyphHeight + 2) * scale; + if (cursor_y + kGlyphHeight * scale > height) { + clear_screen(); + } +} + +void draw_character(char c) { + const uint16_t cell_width = (kGlyphWidth + 1) * scale; + if (cursor_x + cell_width > width) { + newline(); + } + + uint8_t glyph[kGlyphWidth]; + get_glyph(c, glyph); + for (int gx = 0; gx < kGlyphWidth; ++gx) { + for (int gy = 0; gy < kGlyphHeight; ++gy) { + const uint16_t color = + (glyph[gx] & (1U << gy)) ? kForeground : kBackground; + for (int sx = 0; sx < scale; ++sx) { + for (int sy = 0; sy < scale; ++sy) { + const size_t x = cursor_x + gx * scale + sx; + const size_t y = cursor_y + gy * scale + sy; + framebuffer[y * width + x] = color; + } + } + } + } + cursor_x += cell_width; +} + +void write_text(const uint8_t *bytes, size_t len) { + for (size_t i = 0; i < len; ++i) { + const uint8_t byte = bytes[i]; + if (byte == '\n' || byte == '\r') { + if (byte == '\n') { + newline(); + } + } else if (byte >= 32 && byte < 127) { + draw_character(static_cast(byte)); + } else if ((byte & 0xC0) != 0x80) { + draw_character('?'); + } + } +} + +void refresh() { + const esp_err_t err = + esp_lcd_panel_draw_bitmap(panel, 0, 0, width, height, framebuffer); + if (err != ESP_OK) { + ESP_LOGW(TAG, "LCD refresh failed: %s", esp_err_to_name(err)); + } +} + +esp_err_t set_backlight(uint32_t percent) { +#if CONFIG_ESP_BOARD_DEV_LEDC_CTRL_SUPPORT + if (!esp_board_manager_check_name(ESP_BOARD_DEVICE_NAME_LCD_BRIGHTNESS)) { + return ESP_OK; + } + + percent = percent > 100 ? 100 : percent; + ESP_RETURN_ON_ERROR(esp_board_manager_init_device_by_name( + ESP_BOARD_DEVICE_NAME_LCD_BRIGHTNESS), + TAG, "initialize lcd_brightness failed"); + + periph_ledc_handle_t *ledc_handle = nullptr; + ESP_RETURN_ON_ERROR(esp_board_manager_get_device_handle( + ESP_BOARD_DEVICE_NAME_LCD_BRIGHTNESS, + reinterpret_cast(&ledc_handle)), + TAG, "get lcd_brightness handle failed"); + + dev_ledc_ctrl_config_t *brightness_config = nullptr; + ESP_RETURN_ON_ERROR(esp_board_manager_get_device_config( + ESP_BOARD_DEVICE_NAME_LCD_BRIGHTNESS, + reinterpret_cast(&brightness_config)), + TAG, "get lcd_brightness config failed"); + + periph_ledc_config_t *ledc_config = nullptr; + ESP_RETURN_ON_ERROR(esp_board_manager_get_periph_config( + brightness_config->ledc_name, + reinterpret_cast(&ledc_config)), + TAG, "get backlight LEDC config failed"); + + const uint32_t max_duty = (1U << ledc_config->duty_resolution) - 1; + const uint32_t duty = percent * max_duty / 100; + ESP_RETURN_ON_ERROR( + ledc_set_duty(ledc_handle->speed_mode, ledc_handle->channel, duty), TAG, + "set backlight duty failed"); + ESP_RETURN_ON_ERROR( + ledc_update_duty(ledc_handle->speed_mode, ledc_handle->channel), TAG, + "update backlight duty failed"); +#else + (void)percent; +#endif + return ESP_OK; +} + +} // namespace + +esp_err_t board_display_init(void) { +#if !CONFIG_ESP_BOARD_DEV_DISPLAY_LCD_SUPPORT + ESP_LOGE(TAG, "selected Board Manager board has no display_lcd device"); + return ESP_ERR_NOT_SUPPORTED; +#else + ESP_RETURN_ON_FALSE( + esp_board_manager_check_name(ESP_BOARD_DEVICE_NAME_DISPLAY_LCD), + ESP_ERR_NOT_FOUND, TAG, "selected board has no display_lcd device"); + ESP_RETURN_ON_ERROR( + esp_board_manager_init_device_by_name(ESP_BOARD_DEVICE_NAME_DISPLAY_LCD), + TAG, "initialize display_lcd failed"); + + dev_display_lcd_handles_t *display = nullptr; + ESP_RETURN_ON_ERROR( + esp_board_manager_get_device_handle(ESP_BOARD_DEVICE_NAME_DISPLAY_LCD, + reinterpret_cast(&display)), + TAG, "get display_lcd handle failed"); + ESP_RETURN_ON_FALSE(display && display->panel_handle, ESP_ERR_INVALID_STATE, + TAG, "display_lcd returned an invalid panel handle"); + panel = display->panel_handle; + + dev_display_lcd_config_t *config = nullptr; + ESP_RETURN_ON_ERROR( + esp_board_manager_get_device_config(ESP_BOARD_DEVICE_NAME_DISPLAY_LCD, + reinterpret_cast(&config)), + TAG, "get display_lcd config failed"); + width = config->swap_xy ? config->lcd_height : config->lcd_width; + height = config->swap_xy ? config->lcd_width : config->lcd_height; + ESP_RETURN_ON_FALSE(width && height, ESP_ERR_INVALID_SIZE, TAG, + "Board Manager returned an invalid display resolution"); + + scale = width >= 800 ? 3 : (width >= 480 ? 2 : 1); + const size_t framebuffer_size = + static_cast(width) * height * sizeof(uint16_t); + framebuffer = static_cast(heap_caps_aligned_alloc( + 64, framebuffer_size, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT)); + ESP_RETURN_ON_FALSE(framebuffer, ESP_ERR_NO_MEM, TAG, + "allocate %u-byte text framebuffer failed", + static_cast(framebuffer_size)); + + clear_screen(); + refresh(); + ESP_RETURN_ON_ERROR(set_backlight(100), TAG, "turn LCD backlight on failed"); + display_ready = true; + ESP_LOGI(TAG, "Board display: %s (%s), %ux%u, text scale=%u", config->chip, + config->sub_type, width, height, scale); + return ESP_OK; +#endif +} + +void board_display_puts(const uint8_t *bytes, size_t len) { + if (!display_ready || !bytes || len == 0) { + return; + } + write_text(bytes, len); + refresh(); +} + +void board_display_stats(float tokens_per_second, + float milliseconds_per_token) { + if (!display_ready) { + return; + } + + char stats[192]; + const int len = snprintf(stats, sizeof(stats), + "ESP32 PLE TINYLM\n\n" + "28.9M PARAMETERS\n" + "%.2F TOKENS/SECOND\n" + "%.1F MS/TOKEN", + tokens_per_second, milliseconds_per_token); + clear_screen(); + if (len > 0) { + write_text(reinterpret_cast(stats), + static_cast(len < static_cast(sizeof(stats)) + ? len + : sizeof(stats) - 1)); + } + refresh(); +} diff --git a/firmware/esp32_tinystories_idf/main/board_display.h b/firmware/esp32_tinystories_idf/main/board_display.h new file mode 100644 index 0000000..c0060c6 --- /dev/null +++ b/firmware/esp32_tinystories_idf/main/board_display.h @@ -0,0 +1,10 @@ +#pragma once + +#include +#include + +#include "esp_err.h" + +esp_err_t board_display_init(void); +void board_display_puts(const uint8_t *bytes, size_t len); +void board_display_stats(float tokens_per_second, float milliseconds_per_token); diff --git a/firmware/esp32_tinystories_idf/main/esp32_tinystories_idf.cpp b/firmware/esp32_tinystories_idf/main/esp32_tinystories_idf.cpp new file mode 100644 index 0000000..308a9e0 --- /dev/null +++ b/firmware/esp32_tinystories_idf/main/esp32_tinystories_idf.cpp @@ -0,0 +1,427 @@ +// ESP-IDF port of the released TinyStories PLE runtime. Board Manager owns +// board/display initialization, and the model stays mapped from flash. + +#include +#include +#include + +#include "freertos/FreeRTOS.h" +#include "freertos/task.h" + +#include "esp_board_manager.h" +#include "esp_check.h" +#include "esp_heap_caps.h" +#include "esp_log.h" +#include "esp_partition.h" +#include "esp_timer.h" + +#define LLM_INT8_ACT 1 +#define LLM_PROFILE 1 +#define LLM_PROFILE_NOW() esp_timer_get_time() +#include "llm.h" +#include "vocab.h" + +#include "board_display.h" + +namespace { + +constexpr int kPromptIds[] = {433, 447, 259, 405}; // "Once upon a time" +constexpr int kGenerateCount = 200; +constexpr int kParallelRowThreshold = 128; +constexpr size_t kStaticSramBytes = 2 * LLM_Q8_MAX_INPUT; +const char *const TAG = "esp32_tinystories_idf"; + +Model model; +Scratch scratch; + +const uint8_t *model_image; +esp_partition_mmap_handle_t model_mapping; +size_t psram_used; +size_t sram_used; + +TaskHandle_t matvec_worker; +TaskHandle_t inference_task; +const QT *matvec_job_tensor; +const int8_t *matvec_job_input; +float matvec_job_input_scale; +float *matvec_job_output; +int matvec_job_split; + +void *alloc_psram(size_t size) { + void *memory = heap_caps_malloc(size, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT); + if (memory) { + psram_used += size; + } + return memory; +} + +void *alloc_required_psram(size_t size, const char *name) { + void *memory = alloc_psram(size); + if (!memory) { + ESP_LOGE(TAG, "required PSRAM allocation failed: %s (%u bytes)", name, + static_cast(size)); + } + return memory; +} + +void *alloc_required_sram(size_t size, const char *name) { + void *memory = heap_caps_malloc(size, MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT); + if (!memory) { + ESP_LOGE(TAG, "required SRAM allocation failed: %s (%u bytes)", name, + static_cast(size)); + return nullptr; + } + sram_used += size; + return memory; +} + +esp_err_t allocate_sram_buffer(float **buffer, size_t elements, + const char *name) { + *buffer = + static_cast(alloc_required_sram(elements * sizeof(float), name)); + return *buffer ? ESP_OK : ESP_ERR_NO_MEM; +} + +esp_err_t allocate_psram_buffer(float **buffer, size_t elements, + const char *name) { + *buffer = static_cast( + alloc_required_psram(elements * sizeof(float), name)); + return *buffer ? ESP_OK : ESP_ERR_NO_MEM; +} + +void matvec_worker_main(void *) { + while (true) { + ulTaskNotifyTake(pdTRUE, portMAX_DELAY); + matvec_i8_range(matvec_job_tensor, matvec_job_input, matvec_job_input_scale, + matvec_job_output, 0, matvec_job_split); + xTaskNotifyGive(inference_task); + } +} + +void matvec_parallel(const QT *tensor, const float *input, float *output) { + static int8_t quantized_input[LLM_Q8_MAX_INPUT]; + if (!matvec_worker || !tensor->w8 || tensor->rows < kParallelRowThreshold) { + MATVEC(tensor, input, output); + return; + } + + float input_scale; + quantize_act(input, tensor->cols, quantized_input, &input_scale); + matvec_job_tensor = tensor; + matvec_job_input = quantized_input; + matvec_job_input_scale = input_scale; + matvec_job_output = output; + matvec_job_split = tensor->rows / 2; + xTaskNotifyGive(matvec_worker); + matvec_i8_range(tensor, quantized_input, input_scale, output, + matvec_job_split, tensor->rows); + ulTaskNotifyTake(pdTRUE, portMAX_DELAY); +} + +esp_err_t allocate_scratch() { + const Cfg &config = model.c; + const size_t dimension = static_cast(config.dim); + const size_t layers = static_cast(config.n_layers); + const size_t ple_dimension = static_cast(config.ple_dim); + const size_t ffn_dimension = static_cast(config.ffn); + const size_t sequence = static_cast(config.seq_len); + + ESP_RETURN_ON_ERROR(allocate_sram_buffer(&scratch.x, dimension, "x"), TAG, + "allocate x failed"); + ESP_RETURN_ON_ERROR(allocate_sram_buffer( + &scratch.h, + ffn_dimension > dimension ? ffn_dimension : dimension, + "h"), + TAG, "allocate h failed"); + ESP_RETURN_ON_ERROR(allocate_sram_buffer(&scratch.qkv, 3 * dimension, "qkv"), + TAG, "allocate qkv failed"); + ESP_RETURN_ON_ERROR(allocate_sram_buffer(&scratch.att, dimension, "att"), TAG, + "allocate att failed"); + ESP_RETURN_ON_ERROR(allocate_sram_buffer(&scratch.g1, ffn_dimension, "g1"), + TAG, "allocate g1 failed"); + ESP_RETURN_ON_ERROR( + allocate_sram_buffer( + &scratch.g2, + ple_dimension > ffn_dimension ? ple_dimension : ffn_dimension, "g2"), + TAG, "allocate g2 failed"); + ESP_RETURN_ON_ERROR( + allocate_sram_buffer(&scratch.ple, layers * ple_dimension, "ple"), TAG, + "allocate ple failed"); + ESP_RETURN_ON_ERROR( + allocate_sram_buffer(&scratch.tmpP, layers * ple_dimension, "tmpP"), TAG, + "allocate tmpP failed"); + ESP_RETURN_ON_ERROR( + allocate_sram_buffer(&scratch.trow, layers * ple_dimension, "trow"), TAG, + "allocate trow failed"); + ESP_RETURN_ON_ERROR(allocate_sram_buffer(&scratch.scores, sequence, "scores"), + TAG, "allocate scores failed"); + + ESP_RETURN_ON_ERROR( + allocate_psram_buffer(&scratch.logits, + static_cast(model.out_vocab), "logits"), + TAG, "allocate logits failed"); + ESP_RETURN_ON_ERROR(allocate_psram_buffer(&scratch.kcache, + layers * sequence * dimension, + "kcache"), + TAG, "allocate kcache failed"); + ESP_RETURN_ON_ERROR(allocate_psram_buffer(&scratch.vcache, + layers * sequence * dimension, + "vcache"), + TAG, "allocate vcache failed"); + return ESP_OK; +} + +esp_err_t copy_norms_to_sram() { + const int dimension = model.c.dim; + const int layers = model.c.n_layers; + const int ple_dimension = model.c.ple_dim; + const float **vectors[3 * LLM_MAX_LAYERS + 2]; + int sizes[3 * LLM_MAX_LAYERS + 2]; + int vector_count = 0; + + vectors[vector_count] = &model.ple_proj_norm; + sizes[vector_count++] = ple_dimension; + for (int layer = 0; layer < layers; ++layer) { + vectors[vector_count] = &model.attn_norm[layer]; + sizes[vector_count++] = dimension; + vectors[vector_count] = &model.ffn_norm[layer]; + sizes[vector_count++] = dimension; + vectors[vector_count] = &model.ple_norm[layer]; + sizes[vector_count++] = dimension; + } + vectors[vector_count] = &model.out_norm; + sizes[vector_count++] = dimension; + + for (int index = 0; index < vector_count; ++index) { + const size_t bytes = static_cast(sizes[index]) * sizeof(float); + void *destination = alloc_required_sram(bytes, "norm vector"); + ESP_RETURN_ON_FALSE(destination, ESP_ERR_NO_MEM, TAG, + "copy norm vector failed"); + memcpy(destination, *vectors[index], bytes); + *vectors[index] = static_cast(destination); + } + + ESP_LOGI(TAG, "norms -> SRAM: %d vectors", vector_count); + return ESP_OK; +} + +esp_err_t stage_weights_to_psram() { + const int expected = llm_core_stage_count(&model); + int staged = llm_stage_core_int8_alloc(&model, alloc_psram); + ESP_RETURN_ON_FALSE(staged == expected, ESP_ERR_NO_MEM, TAG, + "staged only %d/%d core tensors", staged, expected); + + void *head_buffer = alloc_required_psram( + llm_stage_int8_bytes(&model.out_head), "output head"); + ESP_RETURN_ON_FALSE(head_buffer, ESP_ERR_NO_MEM, TAG, + "stage output head failed"); + llm_stage_int8(&model.out_head, head_buffer); + ++staged; + + ESP_LOGI(TAG, "weights -> PSRAM: %d tensors, %.2f MB allocated", staged, + psram_used / 1048576.0); + return ESP_OK; +} + +void initialize_parallel_matvec() { +#if CONFIG_FREERTOS_NUMBER_OF_CORES > 1 + inference_task = xTaskGetCurrentTaskHandle(); + const int inference_core = xPortGetCoreID(); + const int worker_core = + (inference_core + 1) % CONFIG_FREERTOS_NUMBER_OF_CORES; + const BaseType_t created = + xTaskCreatePinnedToCore(matvec_worker_main, "llm_matvec", 4096, nullptr, + 2, &matvec_worker, worker_core); + if (created == pdPASS) { + model.layer_matvec = matvec_parallel; + model.head_matvec = matvec_parallel; + ESP_LOGI(TAG, "inference core=%d, matvec worker core=%d", inference_core, + worker_core); + } else { + ESP_LOGW(TAG, "matvec worker creation failed; using one core"); + } +#else + ESP_LOGI(TAG, "single-core target; using one-core int8 matvec"); +#endif +} + +uint32_t model_fingerprint() { + uint32_t fingerprint = 2166136261U; + for (size_t index = 0; index < model.image_bytes; ++index) { + fingerprint ^= model_image[index]; + fingerprint *= 16777619U; + } + return fingerprint; +} + +esp_err_t load_model_partition() { + const esp_partition_t *partition = esp_partition_find_first( + ESP_PARTITION_TYPE_DATA, static_cast(0x40), + "model"); + ESP_RETURN_ON_FALSE(partition, ESP_ERR_NOT_FOUND, TAG, + "model partition not found"); + + const void *mapped_model = nullptr; + ESP_RETURN_ON_ERROR(esp_partition_mmap(partition, 0, partition->size, + ESP_PARTITION_MMAP_DATA, &mapped_model, + &model_mapping), + TAG, "map model partition failed"); + model_image = static_cast(mapped_model); + + const int load_result = llm_load(model_image, &model); + ESP_RETURN_ON_FALSE(load_result == 0, ESP_ERR_INVALID_RESPONSE, TAG, + "invalid model image (llm_load=%d)", load_result); + ESP_RETURN_ON_FALSE(model.image_bytes <= partition->size, + ESP_ERR_INVALID_SIZE, TAG, + "model image exceeds its partition"); + ESP_RETURN_ON_FALSE(VOCAB_N == model.out_vocab, ESP_ERR_INVALID_SIZE, TAG, + "tokenizer/model mismatch: vocab.h=%d, model=%d", VOCAB_N, + model.out_vocab); + + const Cfg &config = model.c; + ESP_LOGI(TAG, + "model: Vin=%d Vout=%d D=%d L=%d H=%d F=%d P=%d (image %.2f MB)", + config.vocab, model.out_vocab, config.dim, config.n_layers, + config.n_heads, config.ffn, config.ple_dim, + model.image_bytes / 1048576.0); + return ESP_OK; +} + +esp_err_t initialize_runtime() { + ESP_RETURN_ON_ERROR(allocate_scratch(), TAG, "scratch allocation failed"); + ESP_RETURN_ON_ERROR(copy_norms_to_sram(), TAG, "norm relocation failed"); + ESP_LOGI(TAG, "hot set -> SRAM: %u B dynamic + %u B static = %u B managed", + static_cast(sram_used), + static_cast(kStaticSramBytes), + static_cast(sram_used + kStaticSramBytes)); + + ESP_RETURN_ON_ERROR(stage_weights_to_psram(), TAG, "weight staging failed"); + initialize_parallel_matvec(); + + ESP_LOGI(TAG, "build: bytes=%u fp=%08x sram=%uB psram=%.2fMB", + static_cast(model.image_bytes), + static_cast(model_fingerprint()), + static_cast(sram_used + kStaticSramBytes), + psram_used / 1048576.0); + ESP_LOGI(TAG, "free: SRAM %.0f KB, PSRAM %.2f MB", + heap_caps_get_free_size(MALLOC_CAP_INTERNAL) / 1024.0, + heap_caps_get_free_size(MALLOC_CAP_SPIRAM) / 1048576.0); + return ESP_OK; +} + +void emit_token(int token) { + if (token < 0 || token >= VOCAB_N) { + return; + } + const uint8_t *bytes = VOCAB_BLOB + VOCAB_OFF[token]; + const int length = VOCAB_OFF[token + 1] - VOCAB_OFF[token]; + fwrite(bytes, 1, static_cast(length), stdout); + fflush(stdout); + board_display_puts(bytes, static_cast(length)); +} + +void generate_story() { + printf("\n>>> "); + int position = 0; + int token = 0; + + for (const int prompt_token : kPromptIds) { + token = prompt_token; + emit_token(token); + llm_forward(&model, token, position++, &scratch); + } + + llm_profile_reset(&scratch); + const int64_t start = esp_timer_get_time(); + int64_t decode_microseconds = 0; + int decoded = 0; + + for (int step = 0; step < kGenerateCount && position < model.c.seq_len; + ++step) { + int best = 0; + float best_value = -1e30f; + for (int vocabulary_index = 0; vocabulary_index < model.out_vocab; + ++vocabulary_index) { + if (scratch.logits[vocabulary_index] > best_value) { + best_value = scratch.logits[vocabulary_index]; + best = vocabulary_index; + } + } + + token = best; + emit_token(token); + const int64_t decode_start = esp_timer_get_time(); + llm_forward(&model, token, position++, &scratch); + decode_microseconds += esp_timer_get_time() - decode_start; + ++decoded; + if ((step & 7) == 0) { + vTaskDelay(1); + } + } + + if (decoded == 0 || decode_microseconds == 0) { + ESP_LOGW(TAG, "no tokens decoded"); + return; + } + + const int64_t total_microseconds = esp_timer_get_time() - start; + const float total_tokens_per_second = + decoded * 1e6f / static_cast(total_microseconds); + const float milliseconds_per_token = decode_microseconds / 1000.0f / decoded; + printf("\n\n--- %d tokens in %.2f s ---\n", decoded, + total_microseconds / 1e6); + printf("throughput: %.2f tok/s (%.1f ms/token compute)\n", + total_tokens_per_second, milliseconds_per_token); + + if (scratch.profile.calls) { + const float divisor = static_cast(scratch.profile.calls) * 1000.0f; + printf("profile ms/token: input %.1f | attn %.1f | ffn %.1f | ple %.1f | " + "head %.1f\n", + scratch.profile.input_us / divisor, + scratch.profile.attn_us / divisor, scratch.profile.ffn_us / divisor, + scratch.profile.ple_us / divisor, scratch.profile.head_us / divisor); + } + board_display_stats(decoded * 1e6f / decode_microseconds, + milliseconds_per_token); +} + +void show_error(const char *message) { + board_display_puts(reinterpret_cast(message), + strlen(message)); +} + +} // namespace + +extern "C" void app_main(void) { + ESP_LOGI(TAG, "TinyStories PLE runtime (ESP-IDF + Board Manager)"); + ESP_ERROR_CHECK(esp_board_manager_print_board_info()); + + const size_t psram_total = heap_caps_get_total_size(MALLOC_CAP_SPIRAM); + ESP_LOGI(TAG, "PSRAM available: %.2f MB", psram_total / 1048576.0); + if (psram_total == 0) { + ESP_LOGE(TAG, "PSRAM is required"); + return; + } + + ESP_ERROR_CHECK(board_display_init()); + const esp_err_t model_error = load_model_partition(); + if (model_error != ESP_OK) { + show_error( + "\nMODEL PARTITION INVALID\nFLASH TINYSTORIES MODEL.BIN AND RESTART\n"); + ESP_LOGE(TAG, "model is unavailable: %s", esp_err_to_name(model_error)); + return; + } + + const esp_err_t runtime_error = initialize_runtime(); + if (runtime_error != ESP_OK) { + show_error("\nRUNTIME INITIALIZATION FAILED\nCHECK SRAM AND PSRAM\n"); + ESP_LOGE(TAG, "runtime initialization failed: %s", + esp_err_to_name(runtime_error)); + return; + } + + generate_story(); + while (true) { + vTaskDelay(pdMS_TO_TICKS(10000)); + } +} diff --git a/firmware/esp32_tinystories_idf/main/idf_component.yml b/firmware/esp32_tinystories_idf/main/idf_component.yml new file mode 100644 index 0000000..bdc96a2 --- /dev/null +++ b/firmware/esp32_tinystories_idf/main/idf_component.yml @@ -0,0 +1,10 @@ +version: 0.1.0 +targets: + - esp32p4 + - esp32s3 + - esp32s31 +dependencies: + idf: ">=5.5.0" + espressif/esp_board_manager: + version: "^0.6" + require: public diff --git a/firmware/esp32_tinystories_idf/partitions.csv b/firmware/esp32_tinystories_idf/partitions.csv new file mode 100644 index 0000000..38bdb37 --- /dev/null +++ b/firmware/esp32_tinystories_idf/partitions.csv @@ -0,0 +1,5 @@ +# A 0x10000 table offset accommodates the largest supported bootloader. +# Name, Type, SubType, Offset, Size, Flags +nvs, data, nvs, 0x11000, 0x6000, +factory, app, factory, 0x20000, 0x190000, +model, data, 0x40, 0x1B0000, 0xE50000, diff --git a/firmware/esp32_tinystories_idf/sdkconfig.defaults b/firmware/esp32_tinystories_idf/sdkconfig.defaults new file mode 100644 index 0000000..e0ebacf --- /dev/null +++ b/firmware/esp32_tinystories_idf/sdkconfig.defaults @@ -0,0 +1,13 @@ +# Common performance settings; Board Manager supplies the hardware configuration and sdkconfig.defaults. applies target-specific memory tuning. +CONFIG_COMPILER_OPTIMIZATION_PERF=y +CONFIG_FREERTOS_HZ=1000 + +# The 14,912,348-byte model and the application share a 16 MB flash layout. +CONFIG_PARTITION_TABLE_CUSTOM=y +CONFIG_PARTITION_TABLE_OFFSET=0x10000 +CONFIG_PARTITION_TABLE_CUSTOM_FILENAME="partitions.csv" +CONFIG_PARTITION_TABLE_FILENAME="partitions.csv" +CONFIG_ESP_COREDUMP_ENABLE_TO_NONE=y + +# Keep diagnostics visible on the selected board's default console. +CONFIG_LOG_DEFAULT_LEVEL_INFO=y diff --git a/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32p4 b/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32p4 new file mode 100644 index 0000000..ce5e22f --- /dev/null +++ b/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32p4 @@ -0,0 +1,4 @@ +# Stream model weights through the widest supported PSRAM and cache configuration. +CONFIG_SPIRAM_SPEED_250M=y +CONFIG_CACHE_L2_CACHE_512KB=y +CONFIG_CACHE_L2_CACHE_LINE_128B=y diff --git a/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32s3 b/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32s3 new file mode 100644 index 0000000..fd1c4ca --- /dev/null +++ b/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32s3 @@ -0,0 +1,9 @@ +# Octal PSRAM at 120 MHz is experimental in ESP-IDF, so keep temperature-based timing retuning enabled. +CONFIG_IDF_EXPERIMENTAL_FEATURES=y +CONFIG_SPI_FLASH_HPM_ENA=y +CONFIG_BOOTLOADER_FLASH_DC_AWARE=y +CONFIG_ESPTOOLPY_FLASHFREQ_120M=y +CONFIG_SPIRAM_SPEED_120M=y +CONFIG_SPIRAM_TIMING_TUNING_POINT_VIA_TEMPERATURE_SENSOR=y +CONFIG_ESP32S3_DATA_CACHE_64KB=y +CONFIG_ESP32S3_DATA_CACHE_LINE_64B=y diff --git a/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32s31 b/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32s31 new file mode 100644 index 0000000..93a7d0d --- /dev/null +++ b/firmware/esp32_tinystories_idf/sdkconfig.defaults.esp32s31 @@ -0,0 +1,2 @@ +# Use the highest PSRAM clock supported by ESP32-S31. +CONFIG_SPIRAM_SPEED_250M=y