# rlx-cortexm

Microcontroller backend for RLX — INT8 kernels for ARMv7E-M
(Cortex-M4F / M7), `no_std`, no heap, ~25 KB code.

Reference target: the [nRF52840 USB
Dongle](https://www.nordicsemi.com/Products/Development-hardware/nRF52840-Dongle)
(64 MHz Cortex-M4F, 256 KB RAM, 1 MB flash). The kernels themselves
are plain `i8` / `i32` Rust, so they also run unchanged on host for
parity tests.

This crate sits **one layer below** the rest of the workspace: it does
**not** implement the `Backend` trait. Microcontroller targets don't
have an optimizer / runtime / memory plan; the kernels are called
directly by a hand-written model file (`src/model.rs`).

## What's here

- **`src/quant.rs`** — `QParams { scale, zero_point }`, `requantize`
  (single f32 multiplier — the M4F FPU makes this cheaper than the
  fixed-point dance CMSIS-NN does for M0-class chips), `sat_i8`,
  `dequant`/`quant`.
- **`src/dense.rs`** — INT8 fully-connected
  (`out = requantize(bias + Σ (x-zp)·(w-zp))`).
- **`src/conv2d.rs`** — INT8 NHWC conv2d, padding-aware, weight layout
  `[O,kH,kW,I]`. NHWC chosen to match CMSIS-NN / TFLite Micro and keep
  the inner loop unit-stride across `c_in` — the right shape for
  ARMv7E-M's `SMLAD`/`SMLALD` packed MAC (not emitted yet, layout is
  the prereq).
- **`src/maxpool.rs`** — INT8 NHWC max-pool, no padding.
- **`src/relu.rs`** — INT8 ReLU clamp at zero-point, in place.
- **`src/argmax.rs`** — argmax over an `i8` slice.
- **`src/model.rs`** — hand-wired TinyConv-MNIST forward pass. Two
  ping-pong scratch buffers, no heap.
- **`src/model_weights.rs`** — auto-generated by `tools/train_mnist.py`:
  per-tensor scales, zero points, INT8 weights, i32 biases, plus one
  embedded test image for the e2e test.
- **`src/reference.rs`** — `f32` oracle (`std`-only) used by parity
  tests.
- **`firmware/`** — separate workspace cross-compiling for
  `thumbv7em-none-eabihf`. USB-CDC echo loop: read 784 bytes (one i8
  image) → run `infer` → write 1 byte (predicted digit).
- **`trainer/`** — native-Rust trainer crate (`rlx-cortexm-trainer`,
  binary `train-mnist`). Builds the TinyConv graph in `rlx_ir`,
  autodiffs through `rlx_opt::autodiff::grad_with_loss`, runs SGD on
  `rlx_cpu`, then quantizes + emits `src/model_weights.rs` and
  `tests/data/test_set.bin`. Reaches ~97% test accuracy in ~50s
  (2 epochs on M-class). Replaces the PyTorch path; the latter
  stays under `tools/train_mnist.py` as a reference / data downloader.
- **`tools/`** — Python: `train_mnist.py` (PyTorch → INT8 + Rust source
  — legacy, also handy for populating the torchvision MNIST cache that
  the Rust trainer reads from), `compare_weights.py` (cross-trainer
  bench: parses two `model_weights.rs` files, diffs scales + arrays,
  and runs an INT8 numpy reference over a `test_set.bin` to report
  per-trainer accuracy + agreement), `mnist_client.py` (host → dongle
  exerciser), `setup.py` (cross-platform one-shot installer for
  rustup target, llvm-tools, pyserial, nrfutil).
- **`Justfile`** — single source of truth for dev commands
  (`just setup`, `just train`, `just compare`, `just bench`,
  `just build-fw`, `just flash`, `just client`, `just demo`).
  `just bench` is the three-way regression: trains, runs host-Rust
  kernels, and (if a flashed dongle is attached) the on-device M4F
  sweep. All three paths should produce bit-identical accuracy —
  first measured at 483/500 = 96.60 % across host numpy, host Rust,
  and a real nRF52840.

## Quantization

Per-tensor symmetric (zero_point = 0) for both activations and weights
in this first cut. Per-channel weight scales are a follow-up — the
kernels already take a per-tensor `mult` because the loop is the right
shape to extend to per-channel without restructuring.

Bias is `i32` in accumulator scale (`x_scale · w_scale`) — same
convention TFLite Micro uses. Output requant is a single f32 multiply
(see `quant::requantize`).

Weight bit-widths: `8` (raw i8), `4` (nibble-packed), `2` (ternary,
crumb-packed). Activations are always i8.

## Install

```toml
[dependencies]
rlx-cortexm = { version = "0.1", default-features = false }  # no_std
```

For host-side parity tests: enable the `std` feature.

## Build / test

Host-side (kernels + model parity, no toolchain needed):

```sh
cargo test -p rlx-cortexm --release
```

Five parity tests + single-image e2e + 500-image bulk validation
(skips if `tests/data/test_set.bin` hasn't been generated by
`tools/train_mnist.py`).

Firmware (needs `thumbv7em-none-eabihf`; `python3 tools/setup.py`
installs everything):

```sh
just build-fw    # cross-compile
just size        # llvm-size sections
just flash       # auto-pick SWD (probe-rs) or DFU (nrfutil)
just client 50   # send 50 MNIST images and print accuracy
```

## Status

Production: 96.6% MNIST accuracy validated bit-for-bit across host
numpy / host Rust / on-device M4F (483 / 500 across all three).

## Gotchas

- **Two workspaces.** `firmware/` is a standalone Cargo workspace (its
  own `Cargo.lock`) because the `thumbv7em-none-eabihf` target doesn't
  mix with the host workspace's BLAS/CUDA/Metal deps. Don't add
  `firmware` as a workspace member.

- **`memory.x` assumes the Nordic open USB bootloader is present** —
  application starts at `0x00001000`, leaving the bootloader at
  `0x000F4000`. For SWD-only flashing (no bootloader) bump
  `FLASH ORIGIN = 0x00000000, LENGTH = 1024K`.

- **`model_weights.rs` is auto-generated** — re-run
  `tools/train_mnist.py` whenever the model architecture changes.
  Don't hand-edit; the file leads with that warning for a reason.

- **VID/PID `0x16c0:0x27dd`** is the V-USB / pid.codes test pair.
  Fine for development, but replace with a registered VID before
  shipping anything that ends up on multiple machines.

- **CMSIS-NN-style packed-MAC isn't wired up yet.** The NHWC
  layout means the inner loop is contiguous over `c_in`, so a future
  ARMv7E-M intrinsics path can use `__SMLAD` over `c_in` pairs without
  reshaping. Gate it on `cfg(target_feature = "dsp")`, not on a Cargo
  feature.

- **No `Backend` trait impl.** The optimizer, memory plan, autodiff,
  graph IR — none of it runs on a microcontroller. If you want a model
  here, transcribe its forward pass to `src/model.rs` and dump weights
  with a script like `tools/train_mnist.py`.

## License

GPL-3.0-only.
