AVR: Bump simavr RAM to 63.5K to fix ML-KEM-1024 stack overflow

The test vectors added in the last commit lead to a stack overflow
on AVR for ML-KEM-1024.

This commit ports the AVR setup from mldsa-native, including its
increase in RAM to nearly 64K.

This fixes the AVR test, and also makes future maintenance of
the AVR backend across mlkem-native and mldsa-native simpler.

- nix/avr: RAMEND 0x81FF -> 0xFFFF, EEPROM E2END 0x3FFF -> 0x7FFF
- platform.mk: __stack=0x81FF -> __DATA_REGION_LENGTH__=0xFC00, so .data/.bss
  grow up from 0x0200 with the stack set at runtime
- avr_wrapper.c / init7.S / exec_wrapper.py: place the argc/argv block at the
  top of RAM and set SP just below it, giving the largest possible stack

Keeping this identical (modulo naming) to mldsa-native lets both projects
share one AVR baremetal harness, so fixes port across as a prefix diff.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>

Author: hanno-becker

.github/workflows/baremetal.yml

@@ -17,7 +17,7 @@ jobs:
       matrix:
         target:
          - runner: ubuntu-latest
-           name: 'AVR ATmega128RFR2 (modified for 32K RAM)'
+           name: 'AVR ATmega128RFR2 (modified for 63.5K RAM)'
            makefile: test/baremetal/platform/avr/platform.mk
            nix-shell: cross-avr
            func: true

nix/avr/default.nix

@@ -7,9 +7,9 @@ let
   # Patched simavr with increased RAM and fixed UART output
   simavr-patched = pkgs.simavr.overrideAttrs (oldAttrs: {
     patches = (oldAttrs.patches or [ ]) ++ [
-      ./simavr-32kb-ram.patch
+      ./simavr-64kb-ram.patch
       ./simavr-uart-output-fix.patch
-      ./simavr-16k-eeprom.patch
+      ./simavr-32k-eeprom.patch
       # Exit-code commands (SIMAVR_CMD_EXIT_CODE_*), not yet in a release
       (pkgs.fetchpatch {
         url = "https://github.com/buserror/simavr/commit/c9354b32e057e409c2fbc9454e26db3b3103c26a.patch";

nix/avr/simavr-16k-eeprom.patch nix/avr/simavr-32k-eeprom.patchnix/avr/simavr-16k-eeprom.patch renamed to nix/avr/simavr-32k-eeprom.patch

@@ -6,4 +6,4 @@ index 1234567..abcdefg 100644
 +++ b/simavr/cores/avr/iom128rfr2.h
 @@ -1,1 +1,1 @@
 -#define E2END                           (0xFFF)
-+#define E2END                           (0x3FFF)
++#define E2END                           (0x7FFF)

nix/avr/simavr-32kb-ram.patch nix/avr/simavr-64kb-ram.patchnix/avr/simavr-32kb-ram.patch renamed to nix/avr/simavr-64kb-ram.patch

@@ -8,8 +8,8 @@
  #define RAMSTART                        (0x200)
 -#define RAMSIZE                         (0x4000)
 -#define RAMEND                          (0x41FF)
-+#define RAMSIZE                         (0x8000)
-+#define RAMEND                          (0x81FF)
++#define RAMSIZE                         (0xFE00)
++#define RAMEND                          (0xFFFF)
  #define XRAMSTART                       (0x0000)
  #define XRAMSIZE                        (0x0000)
  #define XRAMEND                         RAMEND

test/baremetal/platform/avr/README.md

@@ -6,4 +6,4 @@ This directory provides a baremetal build and test environment for AVR MCUs, usi
 
 This is primarily a vehicle to test that mlkem-native builds and is functionally correct in 16-bit C implementations. For actual practical use on 16-bit MCUs, stack usage would need to be reduced.
 
-**Note:** We currently need 32K of RAM, more than any MCU supported by `simavr`; we therefore use a patched version of `simavr` where Atmega128rfr2 is given 32K of RAM. To test this, you must work in the `nix .#avr` shell specified in nix flake.
+**Note:** We currently need close to the full 64K data address space of the AVR architecture, more than any MCU supported by `simavr`; we therefore use a patched version of `simavr` where Atmega128rfr2 is given 63.5K of RAM (0x0200-0xFFFF). To test this, you must work in the `nix .#cross-avr` shell specified in nix flake.

test/baremetal/platform/avr/avr_wrapper.c

@@ -14,7 +14,17 @@
 /* Register for sending commands (e.g. exit codes) to simavr */
 AVR_MCU_SIMAVR_COMMAND(&GPIOR0);
 
-#define RAM_BASE 0x2000
+/* The argc/argv block is placed at the top of RAM, just below 16 bytes of
+ * scratch stack used during startup. The exec wrapper chooses the base
+ * address RAM_TOP - blocksize and stores it in the first two bytes of
+ * EEPROM, followed by the block itself. The stack grows downwards from
+ * just below the block, so binaries with little argument data
+ * automatically get the largest possible stack. */
+#define RAM_TOP 0xFFF0
+
+/* Base address of the argc/argv block, read from EEPROM. Used by the
+ * argc/argv register setup in init7.S. */
+uint16_t mlk_argv_base;
 
 static int uart_putchar(char c, FILE *stream)
 {
@@ -27,11 +37,22 @@ static int uart_putchar(char c, FILE *stream)
 /* Set up stdout stream for avr-libc printf */
 static FILE mystdout = FDEV_SETUP_STREAM(uart_putchar, NULL, _FDEV_SETUP_WRITE);
 
-/* Init6 function - copy EEPROM to RAM */
+/* Init6 function - copy argc/argv block from EEPROM to top of RAM and
+ * point the stack just below it. The scratch stack above RAM_TOP is used
+ * for the eeprom_read_* calls; the default stack location (RAMEND of the
+ * unpatched MCU) may lie inside .data/.bss and must not be used. */
 void setup_args(void) __attribute__((naked, section(".init6"), used));
 void setup_args(void)
 {
-  eeprom_read_block((void *)RAM_BASE, (void *)0x0000, 0x4000);
+  uint16_t base;
+  SPH = 0xFF;
+  SPL = 0xFF;
+  base = eeprom_read_word((const uint16_t *)0);
+  eeprom_read_block((void *)base, (const void *)2, (size_t)(RAM_TOP - base));
+  mlk_argv_base = base;
+  base--;
+  SPH = (uint8_t)(base >> 8);
+  SPL = (uint8_t)(base & 0xFF);
 }
 
 /* This is run as part of the init sequence, after setting up the stack

test/baremetal/platform/avr/exec_wrapper.py

@@ -7,6 +7,14 @@
 import os
 import tempfile
 
+# The argc/argv block is placed at the top of RAM, just below 16 bytes of
+# scratch stack used during startup (see avr_wrapper.c). The stack grows
+# downwards from just below the block.
+RAM_TOP = 0xFFF0
+
+# Patched EEPROM size (see nix/avr/simavr-32k-eeprom.patch)
+EEPROM_SIZE = 0x8000
+
 
 def intel_hex_line(addr, data):
     """Generate Intel HEX format line"""
@@ -25,6 +33,10 @@ def intel_hex_line(addr, data):
 def create_eeprom_hex(args, output_file):
     """
     Create EEPROM hex file from command line arguments.
+
+    EEPROM layout: 2 bytes block base address (little-endian), followed by
+    the argc/argv block to be copied to that address:
+    argc (2 bytes) + argv array (len(args) * 2 bytes) + packed strings.
     """
     # First arg should be binary name (strip path)
     args = [os.path.basename(args[0])] + args[1:]
@@ -38,15 +50,23 @@ def create_eeprom_hex(args, output_file):
         strings_data.extend(arg.encode("utf-8"))
         strings_data.append(0x00)  # Null terminator
 
-    # Step 2: Calculate where strings will be in RAM
-    # Layout: argc (2 bytes) + argv array (len(args) * 2 bytes) + strings
+    # Step 2: Calculate where the block will be in RAM
     argc_size = 2
     argv_size = len(args) * 2
-    strings_ram_base = 0x2000 + argc_size + argv_size
+    block_size = argc_size + argv_size + len(strings_data)
+    if block_size + 2 > EEPROM_SIZE:
+        print(
+            f"Error: argument block of {block_size} bytes does not fit in EEPROM",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    base = RAM_TOP - block_size
+    strings_ram_base = base + argc_size + argv_size
 
-    # Step 3: Build data starting with argc
+    # Step 3: Build EEPROM data: block base address, then argc
     data = bytearray()
-    data.extend([len(args) & 0xFF, (len(args) >> 8) & 0xFF])  # argc (little-endian)
+    data.extend([base & 0xFF, (base >> 8) & 0xFF])
+    data.extend([len(args) & 0xFF, (len(args) >> 8) & 0xFF])
 
     # Step 4: Build argv array with pointers to RAM addresses
     for offset in string_offsets:
@@ -96,7 +116,7 @@ def main():
 
         # Run with simavr - enable UART output
         # Note that we use a patched version of simavr where atmega128rfr2
-        # has 32K of RAM. This is purely for testing purposes.
+        # has 63.5K of RAM. This is purely for testing purposes.
         cmd = [
             "simavr",
             "-m",
@@ -126,12 +146,12 @@ def main():
                     filtered_lines.append(clean_line)
             output = "\n".join(filtered_lines)
 
-        # simavr does not propagate the guest exit code (returncode is always
-        # 0), so a guest-side failure would otherwise be reported as success.
-        # As a stopgap until simavr's exit-code mechanism is backported (#1728),
-        # scan the output for "ERROR" and fail instead. This catches both
-        # the UBSan-trap abort() in avr_wrapper.c and the CHECK(...) macros in
-        # the tests, which print "ERROR (file,line)" on failure.
+        # The guest exit code is propagated via simavr's exit-code commands
+        # (see avr_wrapper.c). As an additional safety net, scan the output
+        # for "ERROR" and fail: this catches failures where the firmware
+        # stops without reaching exit(), and both the UBSan-trap abort() in
+        # avr_wrapper.c and the CHECK(...) macros in the tests print
+        # "ERROR (file,line)" on failure.
         #
         # On failure, write to stderr, otherwise stdout.
         if "ERROR" in output:

test/baremetal/platform/avr/init7.S

@@ -7,16 +7,18 @@
 /* AVR calling convention for main(int argc, char** argv):
  * - argc (int, 16-bit) goes in r24:r25 (r24=low, r25=high)
  * - argv (char**, 16-bit pointer) goes in r22:r23 (r22=low, r23=high)
+ *
+ * The argc/argv block was copied to mlk_argv_base by setup_args (init6):
+ * argc at mlk_argv_base, argv array at mlk_argv_base + 2.
  */
 
-/* Load argc from 0x2000 */
-ldi r30, 0x00
-ldi r31, 0x20
+lds r30, mlk_argv_base
+lds r31, mlk_argv_base+1
 ld  r24, Z
 ldi r25, 0
 
-/* Set argv = 0x2002 */
-ldi r22, lo8(0x2002)
-ldi r23, hi8(0x2002)
+movw r22, r30
+subi r22, lo8(-2)
+sbci r23, hi8(-2)
 
 /* Fall through to init8 - no ret instruction */

test/baremetal/platform/avr/platform.mk

@@ -7,11 +7,10 @@ CROSS_PREFIX=avr-
 CC=gcc
 
 # AVR target configuration
-# We would need ATMega256rfr2 with 32K RAM, but it's not supported by simavr.
-# We instead use ATMega128rfr1 with 16K RAM, but modify its specification in
-# simavr to bump the RAM to 32K.
-# Once simavr supports ATMega256rfr2, or once mlkem-native has [an option for]
-# lower stack usage, this should be changed.
+# ML-KEM-1024 (together with the RAM-resident test vectors) needs more than
+# the 16K RAM of any MCU supported by simavr, so we use ATMega128rfr2 with its
+# RAM specification in simavr bumped from 16K to the maximal 63.5K
+# (0x0200-0xFFFF).
 AVR_MCU ?= atmega128rfr2
 AVR_FREQ ?= 16000000UL
 
@@ -35,12 +34,16 @@ CFLAGS += \
 
 CFLAGS += $(CFLAGS_EXTRA)
 
-# Non-standard stack end: 0x81FF = 0x200 + 8K instead of default 0x41FF
+# Memory layout (data address space 0x0200-0xFFFF):
+# - .data/.bss grow upwards from 0x0200 (data region length raised from the
+#   default 16K to the full RAM size)
+# - the argc/argv block sits at the top of RAM, with the stack growing
+#   downwards from just below it (set up at runtime, see avr_wrapper.c)
 LDFLAGS += \
 	-mmcu=$(AVR_MCU) \
 	-Wl,--gc-sections \
 	-Wl,--relax \
-	-Wl,--defsym=__stack=0x81FF \
+	-Wl,--defsym=__DATA_REGION_LENGTH__=0xFC00 \
 	-Wl,--undefined=_simavr_command_register \
 	-Wl,--section-start=.mmcu=0x910000 \
 	-lprintf_min

test/src/test_mlkem.c

@@ -30,9 +30,9 @@
 
 #if !defined(MLK_CONFIG_NO_KEYPAIR_API) && \
     !defined(MLK_CONFIG_NO_ENCAPS_API) && !defined(MLK_CONFIG_NO_DECAPS_API)
-/* Force out-of-line so the ~6.4KB of stack buffers below stay in short-lived
- * frames and don't stack up in main() -- crucial on memory-constrained
- * targets such as AVR (32K RAM). */
+/* Forbid inlining so stack buffers below stay in short-lived
+ * frames and don't stack up in main(). This helps testing on
+ * memory-constrained targets such as AVR. */
 static MLK_NOINLINE int test_keys_core(uint8_t pk[CRYPTO_PUBLICKEYBYTES],
                                        uint8_t sk[CRYPTO_SECRETKEYBYTES],
                                        uint8_t ct[CRYPTO_CIPHERTEXTBYTES],