Codegen: Compiling the Control Loop to a Native Library¶

The codegen pipeline turns a closed-loop control step (estimator + controller on a plant) into a static computation graph that is verified to float-exactness against numpy and then lowered to a native library (.so) via a Zig comptime-unrolled VM.

The motivation is deployment: the shipped .so must provably agree with the Python loop. Rather than reimplementing the control math by hand in a compiled language (and risk subtle drift), the graph is the Python execution, transcribed. The only things that can be wrong are the ~20 primitive op lowerings, which are verified once and reused for every component.

This is an XLA/JAX-style tracing model. Components run once with abstract Tracer values instead of real arrays; every operation they perform is captured as a node in a Graph. The graph is then replayed on real inputs by the interpreter (the correctness oracle) or lowered to Zig by lower_zig.

The pipeline lives in src/shinro/codegen/; the Zig VM lives in src/shinro/runtime/ (see src/shinro/runtime/README.md). The design narrative is in lab-notes/daily/2026-08-24.md.

Pipeline overview¶

                    ┌─────────────────────────────────────────────┐
                    │             shinro.codegen                 │
  component ───────▶│ trace_node()   ──▶ NodeGraph (per component)│
                    │ compose()      ──▶ ComposedGraph (one tick) │
                    │ interpret()    ──▶ numpy arrays             │
                    │ lower_zig()    ──▶ Zig → base.so  (shipped) │
                    └─────────────────────────────────────────────┘

Trace each component once (trace_node). A component is run with Tracer values instead of real arrays; its self.bk is swapped to a TraceBackend, so every bk.* call emits a graph node. Concrete parameters (gain K, matrices A/B) are lifted to const nodes the moment they touch a traced operation — this is the "fixed as compiled" property: shapes and parameters are baked at trace time.
Compose (compose). Stitch the per-component graphs into a single graph representing one tick of the closed loop, wiring them per the fixed ABC dataflow. Shape mismatches (e.g. the KF's (n,1) column vector vs the LQR's (n,) flat) are bridged by auto-inserted reshape nodes — the XLA approach.
Interpret (interpret). Replay the graph on real numpy inputs as a correctness oracle. If the interpreter's output matches a live NumpyBackend run to float-exactness, the tracer is sound.
Lower (lower_zig). Walk the graph and emit Zig — a src/shinro/runtime/ module exposing a shinro_step C-ABI function with baked constants, compiled to a .so. The generated graph is written to src/shinro/runtime/graph_data.zig; the comptime VM that executes it is src/shinro/runtime/lower.zig.

Module map¶

Module	Role
`codegen/tracing.py`	`Tracer` (abstract value), `Graph` / `Node` (graph records), shape checking. Operator overloads (`@`, `+`, `-`, `*`, `.T`) record nodes.
`codegen/trace_backend.py`	`TraceBackend` — a recording `ArrayBackend` that emits nodes for the named `bk.*` methods components call.
`codegen/trace_node.py`	`trace_node` / `trace_node_with_state` — run one component call under a `TraceBackend` and return a `NodeGraph`.
`codegen/infer_contract.py`	Auto-infer a component's I/O contract from its ABC (`Controller` → `compute`, `StateEstimator` → `estimate`) and `inspect.signature`; detect recurrent state via attr-diff.
`codegen/ops.py`	Op-handler registry (`OP_HANDLERS`). The data-driven "switch" — adding a new op is one `@register_op` decorator.
`codegen/interpreter.py`	`interpret` / `interpret_step` — replay a graph on real numpy inputs.
`codegen/compose.py`	`compose` — merge per-component graphs into one closed-loop step graph, auto-inserting `reshape`/`clip`.
`codegen/lower_zig.py`	Emit `src/shinro/runtime/graph_data.zig` (the graph as Zig constants) from a composed graph.
`demo_codegen.py` (repo root)	Runnable demo: traces KF+LQR for the base and cartpole plants, composes, and verifies each stage against a live numpy loop.
`src/shinro/runtime/` (Zig)	`build.zig` (build script), `lower.zig` (comptime-unrolled VM), `linalg.zig` (shared linear-algebra kernels), `graph_data.zig` (generated graph).
`scripts/gen_base.py`	Serializes the `base_tracking` composed graph to `src/shinro/runtime/graph_data.zig` (the `make zig-gen` target).
`scripts/trace_component.py`	Standalone "does my component trace?" gate: inventory of registered components, inferred trace contract (`--list`), and trace + interpret-vs-live oracle check (bit-exact) for any config TOML.

The tracing model¶

`Tracer` — the abstract value¶

A Tracer stands in for an ndarray during tracing. It carries only its concrete shape and its graph node id — no data. Operations on tracers (@, +, -, *, -, .T) emit nodes into the graph and return new tracers.

__array_ufunc__ = None on Tracer is critical: it tells numpy to defer to the tracer's reflected operators (__rmatmul__ etc.) when a numpy array interacts with a tracer, instead of coercing the tracer via np.asarray. Without it, A @ tracer would silently compute with garbage instead of recording a matmul node.

`_lift` — freezing constants¶

When a concrete numpy array (a precomputed gain K, a config-baked matrix) touches a tracer, _lift bakes it into a const node. This is the "fixed as compiled" freezing: the graph carries literal parameter values, so a deployed .so needs no runtime configuration — everything is baked in.

The backend swap¶

trace_node temporarily replaces the component's self.bk with a TraceBackend that records into a fresh Graph. Named ArrayBackend methods (eye, zeros_like, inv, clip, where, copy, ...) emit nodes; bare operators (@, +, ...) are handled by the Tracer overloads. Any method the backend doesn't implement raises a NotImplementedError naming the op to register — a loud, actionable signal, so the op set grows incrementally, driven by real components.

After the call, the original backend and all array instance attrs are restored in a finally, so a traced call never leaves the component polluted with tracers (which would reference dead node ids on the next trace).

Auto-inferred contracts¶

A component needs no per-component tracer metadata. The contract is derived from:

Which method to call — from the ABC the component implements (Controller → compute, StateEstimator → estimate).
Input names and count — from inspect.signature of that method.
Input shapes — supplied by the caller, from the scenario's plant dimensions.
Constants — auto-lifted the moment self.K @ tracer executes.
State — auto-detected at trace time via attr-diff: the tracer snapshots the id() of every array-valued instance attr before the call; any attr whose id() changed after is a recurrent edge (e.g. KalmanFilter.x_hat, KalmanFilter.P).

The only time codegen/ needs an edit for a new component is if it uses a new ArrayBackend op — a one-decorator change in ops.py.

The composition pass¶

compose(estimator, controller, plant_dims, input_limits) wires the fixed ABC dataflow:

y (measurement) ──▶ Estimator ──x_hat──▶ Controller ──u──▶ [clip] ──▶ output
x_ref ────────┬──────────────────────────▶ Controller (reference role)
              └─(regulator: sub)─▶ e = x_hat - x_ref ─▶ Controller (state role)
state_x_hat (recurrent) ─▶ Estimator
state_P (recurrent) ─────▶ Estimator   (any state_* port the trace declares)

Input ports: y, x_ref, u_prev, and every state_* placeholder the traced estimator declares (for the KF: state_x_hat and state_P).
Output ports: u, plus the recurrent state outputs (state_x_hat, state_P, state_u_prev — fed back as inputs next tick). A state attr the trace detected as mutated without a matching pre-injected placeholder raises: that recursion would be silently frozen at its trace-time value (e.g. the KF's covariance collapsing to a one-step gain).
Controller role mapping: the controller's inputs are mapped by role from its compute() signature (_CONTROLLER_INPUT_ROLES in compose.py), not by hardcoded names. Roles: state (names like x0, current_state, state), reference (x_ref, target_state, target), u_prev.
A controller that declares a reference input (LQR, MPPI) receives x_hat and x_ref separately.
A regulator with no reference input (MPC_LTI, MPC_DeltaU) receives the error state e = x_hat - x_ref instead — regulating e to zero tracks x_ref (a sub node bridges the two). Exact for A = I plants; general A needs an (A - I) x_ref feedforward to avoid steady-state offset.
A controller declaring u_prev (MPC_DeltaU) shares the estimator's previous-control recurrent port — the same value feeds both, and state_u_prev closes the loop.
Unmapped input names (e.g. SMC's dynamics terms f_x/g_x, which need a different wiring model) raise at compose time rather than mis-wiring.
Controller recurrent state: the same state_* mechanism applies to the controller side — e.g. PID's _integral/_prev_error thread as state_integral/state_prev_error ports (leading underscores are stripped from port names). A Python branch on instance state would bake the traced path forever, so stateful controllers must express selection as data: PID's first-tick gate is a where on a 0/1 _has_run recurrent port, and its anti-windup back-calculation is an elementwise mask (ne + where), replacing the old if bk.any(...) Python branch.
clip: if input_limits is provided (from [scenario.input_limits]), a clip node is inserted on the controller output.
Auto-reshape: where shapes mismatch (KF (n,1) → LQR (n,)), reshape nodes are inserted. The estimator's (n,1) state output is flattened to the controller's (n,) expectation; (n,) feeds are reshaped to the estimator's (n,1).
_merge_and_rewire: subgraph input nodes are placeholders, not copied — consumers are rewired directly to combined-graph source nodes. Subgraph output nodes are markers and are skipped; compose declares the combined outputs itself.

The wiring is not a per-scenario edge dict — it's the fixed ABC dataflow, the same for every scenario. What's scenario-specific (clip limits, vector dims) comes from the scenario config.

The interpreter¶

interpret(graph, inputs) walks the graph in execution order (the nodes are already topologically sorted — emitted in execution order), dispatches each through OP_HANDLERS[node.op], and collects named outputs. It's a 5-line loop over the registry. If its output matches a live NumpyBackend run to float-exactness, the tracer is sound. The interpreter is the correctness oracle: every test verifies interpret(graph, inputs) against a live numpy computation.

interpret_step is a convenience that splits the outputs into non-state outputs and recurrent state_* outputs (which feed back as inputs next tick).

Ops¶

The op registry lives in ops.py. Register a handler:

@register_op("matmul")
def _matmul(node, values, inputs):
    return values[node.inputs[0]] @ values[node.inputs[1]]

An unsupported op raises NotImplementedError naming the op to add and listing available ops. The current set (from ops.py):

const, input, output, matmul, add, sub, mul, neg, transpose, inv, reshape, clip, where, copy, any, tanh, relu, div, exp, argmax, one_hot, slice.

Lowering to Zig (shipped)¶

The lowerer (codegen/lower_zig.py) walks a composed graph and serializes it to src/shinro/runtime/graph_data.zig — the nodes, const blob, and per-node shapes become Zig compile-time constants. The runtime VM (src/shinro/runtime/lower.zig) is a comptime-unrolled interpreter: one inline for over the graph nodes with a switch (node.op) dispatch, where each node's rows/cols are comptime loop bounds. This mirrors the XLA model of the Python tracer:

Fixed at compile time — shapes, constants, and the op set are baked; there is no heap allocation and no runtime dispatch. Each node's work is statically unrolled.
Static buffers — a single stack array sized from the graph's total buffer footprint (buf_len) is sliced per-node via offsets; no per-op allocation.
Pure dataflow — inputs arrive via an inp slice, outputs are written to an out slice, and there are no side effects.

One deliberate nuance: the no-heap property means "no per-op allocation and no op dispatch at runtime", not "no numeric iteration inside an op". Ops such as inv already do runtime LU iteration inside their comptime-shaped buffer — the same way XLA lowers tf.linalg.inv or Select to runtime loops. This is what keeps the deployment provably correct: the graph shape is known at compile time, even when the numeric work inside an op is data-dependent.

The convergence-iterative solve_qp op follows this same shape. Instead of a comptime-bounded workspace, it drives a statically-allocated OSQP solver generated by osqp.OSQP().codegen(folder, parameters="vectors") into src/shinro/runtime/codegen/emosqp/ (see scripts/gen_emosqp_test.py). The problem data (P, A, l, u) and the pre-factorized KKT matrix are baked into the solver global at generation time; only the linear cost q is updated per tick via osqp_update_data_vec, so there is no per-tick allocation and no libosqp.so dependency. The node's output is the full solution (length n_vars of the baked problem); MPC slices out u[:m] with a downstream slice op.

The generated .so exposes a shinro_step C-ABI function; tests/test_zig_lowering.py loads it with ctypes and cross-checks its output against the Python interpreter to float-exactness (the solve_qp op is exercised by the MPC graph fixture, which compares the codegen solver's output against the interpreter's OSQP solve to within OSQP's tolerance).

Building and testing the Zig layer¶

Requires zig on PATH:

make test-zig    # zig-gen → zig-build → zig test → pytest tests/test_zig_lowering.py

The individual steps are make zig-gen (serialize the base_tracking graph to src/shinro/runtime/graph_data.zig) and make zig-build (compile src/shinro/runtime/build.zig into build/lib/libbase.so). The .so lands in build/ (gitignored).

src/shinro/runtime/build.zig accepts two build options that select which generated graph and which baked OSQP solver a build compiles in, without touching the shared paths:

zig build --build-file src/shinro/runtime/build.zig --prefix build/ \
    -Dgraph=<path-to-graph_data.zig> -Dsolver_dir=<path-to-bake-dir>

-Dgraph defaults to src/shinro/runtime/graph_data.zig; -Dsolver_dir defaults to src/shinro/runtime/codegen/emosqp/. This is how a second MPC bake (e.g. MPC_DeltaU, n_vars=45) coexists with the shipped MPC_LTI one: bake it into a separate directory and build the DeltaU graph against it. The bake's solver_meta.zig (pub const n_vars) feeds a comptime check in src/shinro/runtime/lower.zig that rejects any graph whose .solve_qp node size doesn't match the bake — a cross-config build fails at compile time instead of silently linking a shape-mismatched solver.

The generated graph also carries a has_solve_qp flag. build.zig reads it and links the OSQP bake only for QP graphs, so LQR/PID binaries omit the OSQP C sources and bake metadata entirely.

Compiling an arbitrary scenario (the e2e workflow)¶

The shipped make zig-build path is welded to the KF+LQR base graph. For any other estimator/controller pair — a different robot, a PID instead of LQR, a Luenberger observer instead of a Kalman filter — the two-script e2e pipeline is generic and never touches the shared paths:

make compile SCENARIO=tests/integration/scenarios/base_tracking.toml

scripts/gen_scenario.py (zig-free) — reads the scenario TOML's [controller]/[estimator] configs, [scenario].input_limits, and the [compile] section; traces + composes via the generic shinro.codegen.build_composed_graph; lowers to an isolated build/<name>/graph_data.zig + manifest. The two-pass trace discovers recurrent state by attr-diff, so a new controller's integral or a new estimator's observer state compose with zero per-component declaration.
scripts/build_scenario.py (zig) — pre-flight zig check, build flags from [compile] (CLI > TOML > default), zig build -Dgraph=<abs>, then verifies before stamping: re-gens the graph in-process and byte-compares manifests (integrity), runs the ctypes oracle (shinro_step vs interpret(), tol 1e-12 / 1e-3 for QP), then stamp_deployment + verify_deployment.

The [compile] section is the build spec: n_x/n_u (baked at trace time), optimize (debug/release → ReleaseFast only), target (cross-compile), solver_dir (required for QP graphs). Unknown keys and invalid optimize values are loud errors. Component swaps are TOML edits: change [controller]/[estimator] and re-run make compile — the graph is regenerated from scratch, and the C-ABI port layout (printed by the gen stage, recorded in the manifest) is the only thing the host must re-pack.

Start from src/shinro/configs/scenarios/_template.toml — a commented scenario skeleton with placeholders for your robot's controller, estimator, and [compile] dims. Copy it, fill in the values, and run make compile.

Build manifest (audit trail)¶

Every build writes a deterministic report next to the artifact (<prefix>/lib/libbase.manifest.json) plus a timestamped archive copy (<prefix>/manifests/<UTC>-<graphsha8>.json). The report describes what is inside the .so: build facts (target triple, optimize mode, zig version), provenance (graph/solver paths + sha256s), solver facts (null for solver-free graphs; otherwise baked n_vars, n_cons, eps, config), and the graph content — the ordered node list (dual Python/Zig op names, wiring, shapes, buffer offsets, aux), the C-ABI port layout, buf_len, has_solve_qp, and the .solve_qp n_vars the graph expects. The graph content comes from a <graph>_manifest.json emitted by lower_zig next to graph_data.zig. No timestamps in the report, so identical inputs produce byte-identical reports — diffing two reports shows exactly what changed op-wise, and the archive records when each combination was built.

Zig coverage of the op set¶

src/shinro/runtime/lower.zig handles a subset of the interpreter's ops — the ops actually emitted by the shipped base_tracking graph (names follow the Zig enum in graph_data.zig; cst/inp/out/where_op are the Zig spellings of const/input/output/where):

const, input, output, matmul, add, sub, mul, div, neg, transpose, inv, reshape, clip, where, any, copy, tanh, relu, exp, argmax, one_hot, slice, sin, cos, stack, solve_qp.

Every interpreter op has a VM switch case. solve_qp is special: the interpreter handler solves with the Python osqp (eps=1e-6), while the VM drives the baked codegen static solver (same problem, same tolerance), so both sides agree within OSQP's tolerance. Adding a new interpreter op is a handler in ops.py plus a switch case in src/shinro/runtime/lower.zig and an enum entry in codegen/lower_zig.py.

Graph invariants¶

A captured graph is a flat, topologically ordered list of nodes. Each node holds an op name, input node ids, a concrete shape, and an opaque attrs dict (baked ndarray for const, target shape for reshape, lo/hi for clip, name for input/output). Inputs/outputs are named ports so the interpreter and composition pass can refer to them symbolically.

Writing a new component¶

Add the component as usual (@register_controller("Foo") + from_config + a standard compute(self, current, target) signature) and it traces with zero tracer-side code. If its compute path uses a new ArrayBackend op, register a handler in ops.py.

Running the demo¶

python demo_codegen.py

Four stages:

Trace a KalmanFilter.estimate() alone; show the graph; verify the interpreter matches live numpy.
Compose KF + LQR into one closed-loop step graph (with auto-reshape and clip); verify the composed step matches a live numpy loop.
Swap the estimator (KF → Luenberger) and re-compose, reusing the LQR graph without re-tracing — the modularity proof.
The cartpole system (4-state, 1-input): define the system, build the LQR gain and Kalman filter from the linearized model, trace both, compose, and verify.

Each stage prints PASS / FAIL based on the max abs error vs a live NumpyBackend reference.

Tests¶

tests/test_codegen.py — single-component tracing (KF, LQR), graph-structure assertions, tracer primitive unit tests.
tests/test_codegen_compose.py — composition, KF+LQR composed step vs numpy loop, estimator/controller swap tests.
tests/test_zig_lowering.py — serializes a composed graph to Zig, builds the .so, and cross-checks its shinro_step output against the Python interpreter to float-exactness.

The full lowering path (graph → .so → cross-check) runs with make test-zig; see Building and testing the Zig layer.