pico: add support for flash-based JIT stream

Introduce jit_stream_flash.c common implementation that leverages
(common) flash behavior that can be written from 1 to 0.

Signed-off-by: Paul Guyot <pguyot@kallisys.net>

Author: pguyot

.github/workflows/build-and-test.yaml

@@ -526,6 +526,19 @@ jobs:
         ulimit -c unlimited
         ./tests/test-heap
 
+    - name: "Test: test-jit_stream_flash with valgrind"
+      if: matrix.library-arch == ''
+      working-directory: build
+      run: |
+        ulimit -c unlimited
+        valgrind --error-exitcode=1 ./tests/test-jit_stream_flash
+
+    - name: "Test: test-jit_stream_flash"
+      working-directory: build
+      run: |
+        ulimit -c unlimited
+        ./tests/test-jit_stream_flash
+
     - name: "Test: test-mailbox with valgrind"
       if: matrix.library-arch == ''
       working-directory: build

doc/src/atomvm-internals.md

@@ -137,7 +137,7 @@ Following BEAM, there are two flavors of the emulator: jit and emu, but eventual
 - Native: the VM only runs native code and all code must be precompiled on the desktop using the JIT compiler (which effectively is a AOT or Ahead-of-Time compiler). In this mode, it is not necessary to bundle the jit compiler on the embedded target.
 - Hybrid: the VM can run native code as well as emulated BEAM code and some code is precompiled on the desktop.
 
-JIT is available on some platforms (currently only x86_64 and aarch64) and compiles Erlang bytecode at runtime. Erlang bytecode is never interpreted. EMU is available on all platforms and Erlang bytecode is interpreted.
+JIT is available on some platforms (currently only x86_64, aarch64 and armv6m) and compiles Erlang bytecode at runtime. Erlang bytecode is never interpreted. EMU is available on all platforms and Erlang bytecode is interpreted.
 
 Modules can include precompiled code in a dedicated beam chunk with name 'avmN'. The chunk can contain native code for several architectures, however it may only contain native code for a given version of the native interface. Current version is 1. This native code is executed by the jit-flavor of the emulator as well as the emu flavor if execution of precompiled is enabled.
 
@@ -158,6 +158,27 @@ A backend implementation is required for each architecture. The backend is calle
 
 A stream implementation is responsible for streaming the machine code, especially in the context of low memory. Two implementations currently exist: `jit_stream_binary` that streams assembly code to an Erlang binary, suitable for tests and precompilation on the desktop, and `jit_stream_mmap` that streams assembly code in an `mmap(2)` allocated page, suitable for JIT compilation on Unix.
 
+### Embedded JIT and Native
+
+On embedded devices, Native mode means the code is precompiled on the desktop and executed natively on the device. This currently works on all ARMv6M devices (Pico and STM32).
+
+The default partition scheme on all platforms is optimized for the Emulated VM which is larger than the JIT or Native VM, and for the Emulated atomvmlib (with no native code for estdlib and no jit library) which is smaller than the JIT atomvmlib (that includes native code for estdlib and jit library).
+
+JIT mode means the Erlang bytecode is compiled to native code directly on the device. This actually is possible on Raspberry Pi Pico by using the flash to store the native code. The first time the code is executed, it is compiled and streamed to flash, and for next runs (including at a future boot), the native code is directly executed.
+
+To achive embedded JIT, it is required to flash the device with the JIT compiler for armv6m which is part of the jit library. This library is quite large, so for Pico boards that come with 2MB of flash, it is required to remove jit modules for other backends. It is also required to change the way code is partitioned.
+
+For example, it is possible to have the following offsets defined in `src/platforms/rp2/src/main.c`:
+
+```
+#define LIB_AVM ((void *) 0x10060000)
+#define MAIN_AVM ((void *) 0x101B0000)
+```
+
+To fit in the lib partition, all networking modules should also be removed (the Pico doesn't have any networking capacity).
+
+After the first run, compiled modules in flash are used unless there is a version mismatch or the application avm or the library avm have been updated on the device. AVM packages end with a section called "end" (0x656E64). When the JIT compiler flashes native code, it changes this name to "END" (0x454E44), by effectively clearing 3 bits in the flash, which is possible without erasing any flash block. Any rewrite of these avm packages will overwrite the section names to "end".
+
 ## The Scheduler
 
 In SMP builds, AtomVM runs one scheduler thread per core.  Scheduler threads are actually started on demand.  The number of scheduler threads can be queried with [`erlang:system_info/1`](./apidocs/erlang/estdlib/erlang.md#system_info1) and be modified with [`erlang:system_flag/2`](./apidocs/erlang/estdlib/erlang.md#system_flag2).  All scheduler threads are considered equal and there is no notion of main thread except when shutting down (main thread is shut down last).

src/libAtomVM/module.c

@@ -353,6 +353,13 @@ Module *module_new_from_iff_binary(GlobalContext *global, const void *iff_binary
                 fprintf(stderr, "Native code chunk found but no compatible architecture or variant found\n");
             }
         }
+    } else {
+        ModuleNativeEntryPoint module_entry_point;
+        uint32_t labels;
+        uint16_t version;
+        if (sys_get_cache_native_code(global, mod, &version, &module_entry_point, &labels) && version == JIT_FORMAT_VERSION) {
+            module_set_native_code(mod, labels, module_entry_point);
+        }
     }
 #endif
 

src/libAtomVM/nifs.c

@@ -5633,6 +5633,8 @@ static term nif_code_server_set_native_code(Context *ctx, int argc, term argv[])
     VALIDATE_VALUE(argv[0], term_is_atom);
     VALIDATE_VALUE(argv[1], term_is_integer);
 
+    avm_int_t labels_count = term_to_int(argv[1]);
+
     term module_name = argv[0];
     Module *mod = globalcontext_get_module(ctx->global, term_to_atom_index(module_name));
     if (IS_NULL_PTR(mod)) {
@@ -5646,10 +5648,12 @@ static term nif_code_server_set_native_code(Context *ctx, int argc, term argv[])
 
     SMP_MODULE_LOCK(mod);
     if (mod->native_code == NULL) {
-        module_set_native_code(mod, term_to_int(argv[1]), entry_point);
+        module_set_native_code(mod, labels_count, entry_point);
     }
     SMP_MODULE_UNLOCK(mod);
 
+    sys_set_cache_native_code(ctx->global, mod, JIT_FORMAT_VERSION, entry_point, labels_count);
+
     return OK_ATOM;
 }
 #endif

src/libAtomVM/sys.h

@@ -296,6 +296,33 @@ void sys_free_platform(GlobalContext *global);
  */
 ModuleNativeEntryPoint sys_map_native_code(const uint8_t *native_code, size_t size, size_t offset);
 
+/**
+ * @brief Get the cache (typically on flash) of native code for a given module
+ *
+ * @details If module is found in cache, return a pointer to the entry point.
+ * Only implemented on platforms with JIT. Implementations on flash typically
+ * check if the jit cache is valid (for lib or for app) and use the pointer to
+ * code as a key.
+ * @param global the global context
+ * @param mod module to return the cache native code for
+ * @param version version of the cache entry (for compatibility with the VM)
+ * @param entry_point entry point to the module, if found
+ * @param labels number of labels
+ * @return \c true if the cache entry was found
+ */
+bool sys_get_cache_native_code(GlobalContext *global, Module *mod, uint16_t *version, ModuleNativeEntryPoint *entry_point, uint32_t *labels);
+
+/**
+ * @brief Add native code to cache for a given module
+ *
+ * @param global the global context
+ * @param mod module to add the native code for
+ * @param version version of the native code
+ * @param entry_point entry point to the module
+ * @param labels number of labels
+ */
+void sys_set_cache_native_code(GlobalContext *global, Module *mod, uint16_t version, ModuleNativeEntryPoint entry_point, uint32_t labels);
+
 #ifdef __cplusplus
 }
 #endif

src/platforms/esp32/partitions.csv

@@ -7,6 +7,5 @@
 # Note: if you change the phy_init or app partition offset, make sure to change the offset in Kconfig.projbuild
 nvs,      data, nvs,       0x9000,     0x6000,
 phy_init, data, phy,       0xf000,     0x1000,
-factory,  app,  factory,  0x10000,   0x1C0000,
-boot.avm,  data, phy,     0x1D0000,    0x40000,
-main.avm, data, phy,     0x210000,   0x100000
+factory,  app,  factory,  0x10000,   0x160000,
+main.avm,  data, phy,     0x170000,    0x290000,

src/platforms/generic_unix/lib/sys.c

@@ -853,4 +853,24 @@ ModuleNativeEntryPoint sys_map_native_code(const uint8_t *native_code, size_t si
     return (ModuleNativeEntryPoint) (native_code + offset);
 #endif
 }
+
+bool sys_get_cache_native_code(GlobalContext *global, Module *mod, uint16_t *version, ModuleNativeEntryPoint *entry_point, uint32_t *labels)
+{
+    UNUSED(global);
+    UNUSED(mod);
+    UNUSED(version);
+    UNUSED(entry_point);
+    UNUSED(labels);
+    return false;
+}
+
+void sys_set_cache_native_code(GlobalContext *global, Module *mod, uint16_t version, ModuleNativeEntryPoint entry_point, uint32_t labels)
+{
+    UNUSED(global);
+    UNUSED(mod);
+    UNUSED(version);
+    UNUSED(entry_point);
+    UNUSED(labels);
+}
+
 #endif

src/platforms/rp2/src/CMakeLists.txt

@@ -55,6 +55,10 @@ else()
     target_compile_definitions(AtomVM PRIVATE PICO_STDIO_USB_CONNECT_WAIT_TIMEOUT_MS=20000)
 endif()
 
+if (AVM_DISABLE_SMP)
+    target_compile_definitions(AtomVM PRIVATE PICO_FLASH_ASSUME_CORE1_SAFE)
+endif()
+
 if (AVM_WAIT_BOOTSEL_ON_EXIT)
     target_compile_definitions(AtomVM PRIVATE WAIT_BOOTSEL_ON_EXIT)
 endif()

src/platforms/rp2/src/lib/CMakeLists.txt

@@ -31,7 +31,6 @@ set(HEADER_FILES
 
 set(SOURCE_FILES
     gpiodriver.c
-    jit_stream_flash.c
     networkdriver.c
     otp_crypto_platform.c
     platform_defaultatoms.c
@@ -110,4 +109,16 @@ if (PICO_CYW43_SUPPORTED)
     define_if_function_exists(libAtomVM${PLATFORM_LIB_SUFFIX} gethostname "unistd.h" PRIVATE HAVE_GETHOSTNAME)
 endif()
 
+if (NOT AVM_DISABLE_JIT)
+    target_sources(
+        libAtomVM${PLATFORM_LIB_SUFFIX}
+        PRIVATE
+            jit_stream_flash_platform.c
+            ../../../../libAtomVM/jit_stream_flash.c
+            jit_stream_flash_platform.h
+            ../../../../libAtomVM/jit_stream_flash.h
+    )
+    target_link_options(libAtomVM${PLATFORM_LIB_SUFFIX} PUBLIC "SHELL:-Wl,-u -Wl,jit_stream_flash_get_nif")
+endif()
+
 target_link_options(libAtomVM${PLATFORM_LIB_SUFFIX} PUBLIC "SHELL:-Wl,-u -Wl,gpio_nif -Wl,-u -Wl,otp_crypto_nif")

src/platforms/rp2/src/lib/jit_stream_flash_platform.c

@@ -0,0 +1,117 @@
+/*
+ * This file is part of AtomVM.
+ *
+ * Copyright 2025 by Paul Guyot <pguyot@kallisys.net>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
+ */
+
+#ifndef AVM_NO_JIT
+
+#include "jit_stream_flash.h"
+
+#include <hardware/flash.h>
+#include <pico/flash.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "rp2_sys.h"
+
+// Helper structures for flash_safe_execute
+struct EraseParams
+{
+    uintptr_t addr;
+};
+
+struct WriteParams
+{
+    uintptr_t addr;
+    const uint8_t *data;
+    size_t len;
+};
+
+static void __not_in_flash_func(do_erase_sector)(void *params_ptr)
+{
+    struct EraseParams *params = (struct EraseParams *) params_ptr;
+    flash_range_erase(params->addr - XIP_BASE, FLASH_SECTOR_SIZE);
+}
+
+static void __not_in_flash_func(do_write_page)(void *params_ptr)
+{
+    struct WriteParams *params = (struct WriteParams *) params_ptr;
+    flash_range_program(params->addr - XIP_BASE, params->data, params->len);
+}
+
+struct JSFlashPlatformContext *jit_stream_flash_platform_init(void)
+{
+    return (struct JSFlashPlatformContext *) 1;
+}
+
+void jit_stream_flash_platform_destroy(struct JSFlashPlatformContext *pf_ctx)
+{
+    UNUSED(pf_ctx);
+}
+
+bool jit_stream_flash_platform_erase_sector(struct JSFlashPlatformContext *pf_ctx, uintptr_t addr)
+{
+    UNUSED(pf_ctx);
+
+    struct EraseParams params = {
+        .addr = addr
+    };
+
+    int r = flash_safe_execute(do_erase_sector, &params, UINT32_MAX);
+    if (UNLIKELY(r != PICO_OK)) {
+        fprintf(stderr, "flash_safe_execute (erase) failed with error %d\n", r);
+        return false;
+    }
+
+    return true;
+}
+
+bool jit_stream_flash_platform_write_page(struct JSFlashPlatformContext *pf_ctx, uintptr_t addr, const uint8_t *data)
+{
+    UNUSED(pf_ctx);
+
+    struct WriteParams params = {
+        .addr = addr,
+        .data = data,
+        .len = FLASH_PAGE_SIZE
+    };
+
+    int r = flash_safe_execute(do_write_page, &params, UINT32_MAX);
+    if (UNLIKELY(r != PICO_OK)) {
+        fprintf(stderr, "flash_safe_execute (write) failed with error %d\n", r);
+        return false;
+    }
+
+    return true;
+}
+
+uintptr_t jit_stream_flash_platform_ptr_to_executable(uintptr_t addr)
+{
+    // Set Thumb bit
+    return addr | 0x1;
+}
+
+uintptr_t jit_stream_flash_platform_executable_to_ptr(uintptr_t addr)
+{
+    // Clear Thumb bit
+    return addr & ~0x1UL;
+}
+
+REGISTER_NIF_COLLECTION(jit_stream_flash, jit_stream_flash_init, NULL, jit_stream_flash_get_nif)
+
+#endif // AVM_NO_JIT