Support for onnx MLPs - #25
Merged
Merged
Conversation
adilfaisal01
commented
Sep 18, 2026
adilfaisal01
commented
Member
- removed onnxruntime as a dependency, instead using onnx
- onnx load graph exposes onnx proto files as a DAG, which can be lowered to zig and .so binary allowing C-ABI contract
- optimized lower.zig, removing mentions of inline for in the inner loop, allowing LLVM IR optimizations, providing significant speed in build time and reductions in binary sizes
- prevented overflow by allowing size to be global rather than local
- comparison harness for ticks
- currently only MLP with a subset of ops are supported, more will be coming soon
The ONNX→Zig policy path reads the graph with `onnx` at compile time and deploys a dependency-free `.so`, so the extra no longer needs the onnxruntime inference engine. `onnx` also joins the dev extra, which fixes the ONNX adapter tests silently skipping under a stock `.[onnx-rl]` install (they guarded on both packages, but only onnxruntime was declared). The checked-in adapter still lazy-imports onnxruntime on its eager path; that path is replaced in a following commit.
Add shinro.codegen.onnx_import: translate onnx.load(...).graph directly into a memoryless ComposedGraph, bypassing the tracer entirely (ONNX is already a dataflow graph, so there is nothing to intercept). Supports Gemm (alpha/beta/ transB), MatMul and Add, Relu/Tanh, and composes Sigmoid from exp/neg/add/div so no new VM op is needed. The observation encoder (integer index selection, mean/std normalization, clipping) is folded in as arithmetic on baked constants, leaving the raw plant state as the only input port. Every emitted node is evaluated eagerly with the ops registry's numpy handler and that result's shape becomes the node's declared shape, so interpreter and lowering semantics cannot drift. Nodes that do not feed the declared output are ignored; unsupported ops/attributes, multi-input models, and batched outputs fail loudly with actionable messages. Verified: 24 unit tests (pytest tests/unit/test_onnx_import.py), ruff clean, pyrefly 0 errors.
Add the action-space surface to the ONNX importer: continuous applies scale/bias, discrete emits an argmax one-hot, and stochastic splits [mean; log_std] and either returns the mean or adds exp(clip(log_std, -10, 2)) * epsilon — mirroring the old runtime _postprocess. A non-deterministic discrete/stochastic policy gains an `epsilon` C-ABI input port: the host supplies the noise (Gumbel for discrete, standard normal for stochastic) and the kernel does only the arithmetic, matching MPPI's port and keeping RNG on the host. Discrete sampling uses the Gumbel-max trick, argmax(logits + g) + one_hot, so no softmax/max VM op is needed. Single-sided action clips are rejected: the lowerer writes floats as Zig hex literals and `inf` is not a Zig identifier, so a bound defaulted to ±inf would fail the build. Gemm's alpha/beta multipliers and the action scale/bias are emitted unconditionally (defaults included) for uniform lowering. Verified: 38 importer unit tests, including a 2000-draw Gumbel-max vs softmax distribution check; all six action-space variants lower via lower_zig into isolated paths; ruff + pyrefly clean.
Rewrite the onnx_rl adapter on top of the ONNX→shinro importer and drop
onnxruntime entirely. Two interchangeable backends: `model_path` imports the
ONNX graph and runs it with the pure-numpy interpreter (eager, no build step),
`artifact_dir` dlopens the scenario's compiled kernel under
lib/lib_neural_network.so and drives it through the shinro_step C ABI, reading
the port layout from the graph manifest so the .onnx file is not needed at
deploy time.
Sampling noise (Gumbel for discrete, standard normal for stochastic) is drawn
host-side from a seeded RNG and fed to the epsilon port, keeping the kernel
pure arithmetic; reset() makes a run reproducible.
Add the frozen OnnxRLConfig dataclass (silences the registry's missing-Config
warning for onnx_rl), strict-parse the config, and reject unknown observation
keys so a typo cannot silently drop normalization.
Behavior changes vs the runtime implementation: actions are float64 (the graph
is f64 throughout), single-sided action_clip is rejected (a ±inf bound cannot
be lowered), normalize without stats raises ValueError, and
observation.add_batch_dim is a no-op (batch-1 is implicit in the 1-D ports).
Verified: 81 unit tests across tests/unit/test_onnx_{rl_adapter,import}.py;
full suite 1225 passed / 5 skipped; ruff + pyrefly clean.
Make `make compile` reach a policy-only scenario end to end: - scenario_gen: a scenario may omit [estimator] when its controller is onnx_rl (any other type still requires one). Such a scenario bypasses build_composed_graph and lowers the importer's graph directly, with the .onnx file's sha256 pinned in the manifest provenance. - [compile].artifact_name selects the installed kernel stem. build.zig installs through an explicit sub-path (Zig's addLibrary would otherwise prefix "lib", turning lib_neural_network into liblib_neural_network); oracle.load_so, stamp, scenario_build and cli thread the name through. The default "libbase" keeps every existing artifact and test byte-identical. - build.zig's build-time readFile cap was 1 MiB, so any graph past ~50k baked constants failed with a misleading "graph_data.zig has no has_solve_qp flag; regenerate it" panic. Raised to 256 MiB. Add a committed toy policy so the ONNX path has a stable subject: tests/fixtures/models/toy_mlp.onnx (26-param 3->4->2 tanh MLP, generated by scripts/gen_toy_onnx.py), its controller config, and a policy-only scenario. The adapter/import tests and the compile e2e use it, a drift guard regenerates it through the script's CLI and compares the graph, and the scenario template documents the new [compile] key (the template drift guard forces that). Add scripts/measure_onnx_policy_scale.py — the cold-build sweep of parameter count against artifact size, graph source size, VM buffers and compile time. Its numbers and the comptime scaling wall (the VM is a ~10^5-parameter design; >=830k params trips @setEvalBranchQuota) are written up in the lab note. Verified: 100 onnx/compile tests including every zig-gated e2e; make test-zig 77 passed / 2 skipped; make test 1233 passed / 5 skipped; ruff + pyrefly clean.
TestOnnxPolicyOracle imports the committed toy policy (tests/fixtures/models/ toy_mlp.onnx) and compiles its four baked variants — continuous, discrete-deterministic, discrete-sampling and stochastic-sampling — into tmp-path kernels, then asserts .so == interpret() on seeded random states (exact for the continuous/one-hot paths, <=1e-14 for the noise paths) plus the closed form, numpy argmax, and the Gumbel/stochastic formulas. This is the first oracle whose subject is a learned policy rather than a hand-written control law, and it proves the importer's output — baked encoder, transposed Gemms, composed activations, action post-processing — compiles bit-for-bit. The shared src/shinro/runtime/graph_data.zig is never touched. Verified: make test-zig 81 passed / 2 skipped; make test 1241 passed / 5 skipped; make lint clean.
configs/controllers/onnx_rl.toml now explains the eager (model_path) and compiled (artifact_dir) backends, the policy-only `make compile` command, the epsilon / host-noise contract, and that add_batch_dim is legacy; the stale `--controller onnx_rl` demo line is removed (no such base config exists). docs/components.md records that onnx_rl is the only controller that may omit [estimator] and the only one with a compiled deployment path. The lab note carries the step 6-7 write-up and the ONNX shift's scale-sweep results. Verified: the updated config strict-parses into OnnxRLConfig; ruff + pyrefly clean.
runtime/lower.zig unrolled both the outer node table and every per-element loop with `inline for` (38 sites, 0 runtime `for`). The outer unroll is cheap (~1300 nodes); the inner ones emit one statement per element, so shinro_step became a single function with ~buf_len statements (100k-340k). Zig's comptime evaluator and LLVM then processed that giant straight-line function: MPPI production builds at N=200/K=15 took 411-1306 s and produced 1.9-6.1 MiB .so files whose .text was ~2x buf_bytes. Keep the outer `inline for (g.nodes)` (op tag and shapes stay comptime, buffer offsets stay comptime constants) and make all 34 inner element loops runtime `for` loops over those comptime-known sizes. ew2/bcast_flat stay inline so shapes and the op tag constant-fold. Result (MPPI, N=200/K=15, ReleaseFast, real plants): - 5.2-11.5x smaller .so - 35-106x faster builds (8-kernel matrix ~90 s total) - 1.4-2.1x faster ticks (the unrolled function was thrashing I-cache) ONNX MLP scale sweep: the >=830k-param comptime branch-quota wall is gone; all six archs (up to 1.88M params) now compile and oracle-verify. Verified: make test 1241 passed/5 skipped (unchanged), make test-zig 81/2, make lint clean, all SMC/MPPI/ONNX and closed-loop oracles pass. Shipped graph_data.zig untouched.
shinro_step declared its per-node scratch as a stack local, `var buf: [g.buf_len]f64`. For large graphs that reserves buf_bytes on the caller's stack — 20.8 MiB at ~1.35M parameters — overflowing the default 16 MiB stack when the kernel is called. The compiler was fine; the call segfaulted. Move it to a file-scope `workspace` in .bss, sized at compile time from g.buf_len and addressed with the same comptime offsets (the helpers keep their `buf` pointer parameter). The graph writes every slot before it is read, so no initialisation is needed. Documented as not reentrant / thread-safe — one caller, one tick at a time — matching the QP path's statically-allocated solver global. Result: a 1.35M-param policy (20.8 MiB workspace) builds and oracle-verifies under the default 16 MiB stack; its .so is ~15 KB of .text plus the .bss workspace, still with no dynamic allocation. Tick latency unchanged within measurement noise. Verified: make test-zig 81 passed / 2 skipped, make test 1241 passed / 5 skipped.
scripts/bench_tick.py dlopens a compiled kernel, drives shinro_step with seeded inputs and recurrent state feedback, and reports ns/tick min/median/p99. It is the A/B yardstick for the VM lowering changes measured in the lab notes.
MD060 table-separator spacing; content unchanged.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.