pb.nu

packages/onnx/src/pb.nu — the protobuf seam between ONNX and the stdlib.

The wire format itself lives in stdlib/ext/protobuf.nu now: varints, tags, length-delimited regions, packed repeated fields, groups, bounds and depth limits, all checked. This file is the adapter that lets model.nu keep reading a message as a flat loop over fields.

A PReader is a checked stdlib reader plus a STICKY error. Every read below is infallible at the call site and returns a neutral value once the reader has failed, so a parser stays a straight while more { ... } loop with no per-field error plumbing — but nothing is silently mis-read: the first failure latches, every later read is a no-op, the enclosing loop stops on pb_more, and pb_err reports what went wrong and where. A sub-message reader carries its failure back to its parent through pb_absorb.

Also here: turning a little-endian byte block into a host buffer of f32 or int64 values. That part was never protobuf — it is what TensorProto.raw_data and the tools' plain .f32 files both need.

API

: PReader { ProtoReader r b failed ProtoError err }

@ pb_empty ProtoError e → PReader

A reader that has already failed — the neutral value a construction failure degrades to, so pb_new needs no error branch at its call site.

@ pb_new ( Vec u ) bytes → PReader

@ pb_failed PReader p → b

@ pb_err PReader p → ProtoError

@ pb_absorb inout PReader p PReader sub → v

Carry a sub-reader's failure into its parent.

@ pb_more inout PReader p → b

More fields to read? False once the reader has failed, so every message loop stops at the first bad field instead of spinning.

@ pb_tag inout PReader p → ProtoTag

Next field tag. Returns field 0 / wire 0 once failed.

@ pb_varint inout PReader p → i

A varint field as i64 (int32/int64/uint/bool/enum all arrive here).

@ pb_i32 inout PReader p → i

A fixed 32-bit field as its raw bit pattern (float bits / fixed32).

@ pb_string inout PReader p → String

A length-delimited field as an owned String. ONNX name/op_type fields are proto3 string, but a STRING attribute's value is bytes and is not required to be UTF-8, so this does NOT validate — it takes the bytes as they are, which is what the previous reader did.

@ pb_bytes inout PReader p → ( Slice u )

A length-delimited field's raw payload, borrowed in place. The slice borrows the caller's model buffer — valid as long as that buffer is.

@ pb_submsg inout PReader p → PReader

An embedded message as its own bounded reader.

@ pb_packed inout PReader p → PReader

A packed repeated field as a bounded reader over its payload.

@ pb_skip inout PReader p ProtoTag tag → v

Skip a field whose tag was already read (groups included).

& c @ nurl_poke_i32 *u base i idx i32 val → v

@ slice_f32_into ( Slice u ) s * u dst i n → v

n little-endian f32 values from s into dst (raw *u, 4-byte stride), writing each one's exact 32-bit pattern — no float round-trip. Elements past the end of the slice read as zero.

@ slice_i64_into ( Slice u ) s * u dst i n → v

n little-endian int64 values from s into dst (8-byte stride).

@ vec_f32_into ( Vec u ) bytes * u dst i n → v

The same two over a whole byte vector — how the tools read a plain little-endian .f32 file.

@ vec_i64_into ( Vec u ) bytes * u dst i n → v

---

tensor_bridge.nu

onnx/tensor_bridge.nu — zero-copy seam between the tensor package's device tensors (DTensor) and the onnx graph runtime (M4a).

The gpu package opens the device's PRIMARY CUDA context, so every package in the process shares one device address space: a DTensor's GkBuf and an Engine's RTensor are pointers into the same heap. That makes composition free of host roundtrips:

preprocess with dtensor_ ops → run the graph → postprocess with dtensor_ ops — the activations never touch the host.

Ownership rules:

• rt_run_dtensor BORROWS the input DTensor (the engine registers the

pointer but never releases it — the caller keeps the DTensor alive until the run's outputs are consumed or copied out).

• dtensor_from_output COPIES the output device-to-device into a fresh

owned DTensor (released with its last owner). A copy (not an ownership transfer) because graph outputs may ALIAS an engine-owned base allocation at an offset (Reshape/Split), and the engine lets its own allocations go at the next rt_reset anyway.

Only TE_F32 tensors cross the seam — graphs compute in float32.

API

@ rt_run_dtensor Engine e__h OGraph g DTensor d → RTensor

Run a graph with a device-resident input: no host staging, no upload. The DTensor must be TE_F32 and shaped as the graph expects. Returns the graph's first output (engine-owned, valid until the next rt_reset).

@ dtensor_from_output GpuKit kit Engine e RTensor t → DTensor

Wrap a run's output as an OWNED device tensor (device-to-device copy — see the ownership note above). The result lives independently of the engine: rt_reset / the engine's last owner do not touch it.

---

model.nu

packages/onnx/src/model.nu — the ONNX schema, over the stdlib codec.

Decodes the subset of the ONNX protobuf needed for feed-forward inference: ModelProto → GraphProto → { NodeProto[], TensorProto initializers, graph input/output names }. Field numbers are the stable ONNX wire layout. Weights (TensorProto.raw_data, little-endian f32) are borrowed in place out of the model buffer and read straight into a host buffer ready to upload to the GPU.

The wire format is stdlib/ext/protobuf.nu, reached through pb.nu — so a malformed model is reported rather than decoded past. Reading a field stays infallible here because a PReader latches its first error; the enclosing loop then stops and the reason reaches onnx_parse_checked.

API

: OAttr { String name i kind f f i i String s ( Vec i ) ints }

── in-memory graph ─────────────────────────────────────────────── One attribute. Captures the value forms used by the ops we run: f (FLOAT, e.g. alpha/epsilon), i (INT, e.g. group), s (STRING, e.g. auto_pad), and ints (INTS, e.g. kernel_shape/strides). kind is unused (the consumer asks by name + form via node_attr_*).

: ONode

: ONode {
    String op_type
    ( Vec String ) inputs
    ( Vec String ) outputs
    ( Vec OAttr ) attrs
}

: OTensor

: OTensor {
    String name
    ( Vec i ) dims
    i nelem
    i dtype  // ONNX DataType: 1=FLOAT, 7=INT64
    ( Vec u ) host  // the values, little-endian (empty = no data)
}

A tensor: name, shape, element count, ONNX data_type, and host data — f32 (4-byte) for FLOAT(1) weights, or i64 (8-byte) for INT64(7) shape / size / anchor tensors, as bytes. host is empty for a graph value placeholder. Every field is owned: an OTensor (and the OGraph holding it) is released with its owner.

@ otensor_host_ptr OTensor t → i

The host data's address (0 when there is none) — for an upload or a peek; valid while the tensor is.

: OGraph

: OGraph {
    ( Vec ONode ) nodes
    ( Vec OTensor ) inits
    String input_name
    String output_name
    String output1_name
}

@ streq s a s b → b

String equality as a bool (nurl_str_eq returns i).

@ node_attr_f ONode n s name f dflt → f

── attribute access ──────────────────────────────────────────────

@ node_attr_i ONode n s name i dflt → i

@ node_attr_int_at ONode n s name i k i dflt → i

The k-th element of an INTS attribute (dflt if attr/element absent).

@ node_attr_ints_len ONode n s name → i

@ node_attr_s ONode n s name s dflt → s

A STRING attribute's value (dflt if absent/empty).

@ onnx_empty_graph → OGraph

An empty graph — what a failed parse returns, and the value every caller initialises its own ~ OGraph binding with.

@ onnx_parse_checked ( Vec u ) bytes → !OGraph ProtoError

Parse a ModelProto, reporting the first wire-format error instead of decoding past it. bytes is borrowed for the call only: every weight block is copied into its own host buffer before this returns, so the caller may free the model buffer as soon as it has the graph.

@ onnx_parse ( Vec u ) bytes → OGraph

The historical entry point: a malformed model yields an empty graph rather than an error. Every existing caller checks the graph it gets back (an empty one runs no nodes), so this keeps their contract — onnx_parse_checked is the one to call when the reason matters.

@ graph_find_init OGraph g s name → i

── lookups ───────────────────────────────────────────────────────

@ graph_free sink OGraph g → v

Let go of g now rather than at the end of its owner's scope. An OGraph owns everything it holds — every node (op string, input / output names, attributes), every initializer (name, dims, host data) and the graph's names — and releases them with its last owner.

---

static_kernels.nu

packages/onnx/src/static_kernels.nu — kernels_static.c for the gpu package's STATIC backend (backend 2: precompiled kernels, no NVRTC, no host compiler, no dlopen — the backend of wasm32 and sealed native builds), derived from the executor.

The kernel set is not listed anywhere. rt_kernel_census (runtime.nu) issues every gkd_* call the executor makes on a gpukit census kit, which records each kernel's entry name and exact source from gpukit's own builders, taking the branches the static backend takes; this file turns the record into C. A kernel the executor gains, or a kernel body gpukit changes, reaches the static set with no edit here. (Up to 0.9.0 the generator mirrored a kernel list by hand, and when the executor moved onto gpukit the mirror kept importing a file that no longer existed: the static and wasm builds were broken for three releases.)

tools/gen_static_kernels.nu is the command-line front end; tests/census_test.nu checks the census covers the executor.

API

: Override { s name s sig s body }

---

main.nu

packages/onnx/src/main.nu — run an ONNX model on the GPU.

onnx <model.onnx> <input.f32> [expected.f32]

Parses the model (pure-NURL protobuf), uploads weights + input to the GPU, executes the graph (Gemm/Relu kernels via packages/gpu + NVRTC), and prints the output. With an expected.f32 it also checks the result against that reference (e.g. an onnxruntime dump) and sets the exit code.

Build: nurlpkg install onnx (or, from the package root, NURL_STDLIB=<repo> ../../nurl.sh src/main.nu)

API

& c @ nurl_peek_f32 *u base i idx → f

@ load_f32 s path → ( Vec u )

Load a raw little-endian f32 file: its values as bytes (4 per element), empty when the file cannot be read.

@ print_f f x → v

@ main → i

---

runtime.nu

packages/onnx/src/runtime.nu — the graph executor.

Walks the ONNX graph in node order (ONNX guarantees topological order), keeping every intermediate tensor resident on the GPU. Initializers and the input are uploaded once; each node dispatches to a gpukit dev-layer kernel (gkd_* — the SAME dtype-generic kernel library the tensor package's DTensor uses; onnx carries no kernel sources of its own); the named output is downloaded at the end. A value map (name → device tensor) threads activations between nodes. Kernels compile lazily and are cached by gpukit (in-process by name, on disk by source hash), and launches chain on the stream with ONE device sync at the end of the walk.

Tensors are N-D (shape vector); the dense path reads dims 0,1 as M,K and the conv path reads NCHW from dims 1,2,3 (batch N is assumed 1).

Memory: an Engine is a handle (an rcbox) — every copy is the same engine, and its last owner releases every device block it allocated and the device; rt_reset lets a run's blocks go early. Nothing is freed by hand (rt_close is an optional early release). The graph a run is given is BORROWED for the engine's use; the caller keeps it.

API

: RTensor { String name i dptr ( Vec i ) shape i nelem }

A device-resident tensor: name, CUdeviceptr (i64), shape, element count.

: EngineImpl

: EngineImpl {
    Gpu g
    GpuKit kit
    ( Vec RTensor ) vals
    i inits_ref
    String input_name
    String output_name
    String output1_name
    b ok
    ( Vec GpuBuffer ) owned
}

The engine's state. g is the kit's device (another owner of it); owned holds every device block a run allocated (rt_reset / the engine's last owner release them). The graph of the current run stays the CALLER's: the engine keeps copies of the three names it needs after the run, and reaches the initializers through inits_ref — the graph's inits Vec as its raw handle word, a view valid while the run is (0 when no run has installed a graph). Holding no part of the graph is what lets a caller run the same graph again, on this engine or another.

: Engine { s ctl }

An Engine is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same engine, and the last owner releases it — every device block it allocated, and the device.

@ Engine_share Engine h → Engine

@ Engine_drop sink Engine h → v

@ _rt_engine_ptr Engine h → *EngineImpl

The state, for this package's other files (tensor_bridge.nu).

@ streq2 s a s b → b

@ ceil_div i a i b → i

& c @ nurl_peek_f32 *u base i idx → f

── initializer metadata (int64 shape/size tensors stay host-side) ──

@ rt_dim RTensor t i ax → i

@ rt_open i ordinal → Engine

Open a device. Kernels compile lazily on first use (gpukit caches them in-process by name and on disk by source hash and architecture).

@ _rt_set_graph * EngineImpl e OGraph g → v

Install the caller's graph for a run: its names are copied, its initializers viewed (see EngineImpl). The graph stays the caller's.

@ rt_own * EngineImpl e GpuBuffer b → i

Keep a device allocation until the next rt_reset (or the engine's last owner) and hand back its address. Aliased tensors (Reshape/Split share or offset an existing buffer) are NOT recorded — only the real gpu_alloc blocks, so each one is released exactly once.

@ rt_reset Engine e__h → v

Release every device buffer allocated during the previous run and clear the value map. Lets one Engine serve many forward passes (e.g. one text prompt per call) without holding the last pass's weights + activations.

@ rt_none → Engine

An engine that is not there (rt_ok F) — the placeholder for "not opened (yet)".

@ rt_ok Engine e__h → b

@ rt_name Engine e__h → s

@ rt_put * EngineImpl e s name i dptr sink ( Vec i ) shape → v

Register a device tensor under name with an explicit shape vector.

@ rt_find * EngineImpl e s name → i

@ rt_at * EngineImpl e i idx → RTensor

: i RT_HOSTI - 0 1

── host-side INT64 tensors ─────────────────────────────────────── The shape-arithmetic chains a torch export leaves behind — Shape → Gather/Unsqueeze/Concat/Cast/Slice → a Resize sizes input — run on the HOST: with a static input shape they are constants, and three-element i64 tensors have no business on the device. A host-int tensor lives in the SAME value map as an RTensor whose dptr is the RT_HOSTI sentinel and whose VALUES sit in the shape vector, so rt_reset frees it like any other entry and the ready gate's rt_find keeps working unchanged.

@ rt_load_inits * EngineImpl e OGraph g → v

@ rt_alloc_out * EngineImpl e s name sink ( Vec i ) shape → i

Allocate a fresh device tensor with shape, register under name, return its dptr.

@ rt_gemm * EngineImpl e ONode n → v

── op handlers ───────────────────────────────────────────────────

@ rt_relu * EngineImpl e ONode n → v

@ rt_conv * EngineImpl e ONode n → v

@ rt_convtranspose * EngineImpl e ONode n → v

Transposed convolution (ConvTranspose). Weight is [Cin, Cout, kh, kw]. Output size per ONNX: O = stride·(I−1) + output_padding + (k−1)·dil+1 − pad_begin − pad_end. The seg Proto upsample is k2/s2/p0 → O = 2·I.

@ rt_maxpool * EngineImpl e ONode n → v

@ rt_batchnorm * EngineImpl e ONode n → v

@ rt_leakyrelu * EngineImpl e ONode n → v

@ rt_ndim RTensor t → i

@ rt_last RTensor t → i

@ rt_binop * EngineImpl e ONode n i op → v

Binary op: 0=Mul 1=Add 2=Sub 3=Div. For commutative ops (Mul/Add) the larger input is the data and the smaller broadcasts (ONNX allows either order); Sub/Div keep in0 op in1. Broadcast mode is inferred from the operand element count: scalar, full, per-inner (a per-anchor stride vector), or per-channel.

: i RT_EW_OPS 4

The broadcast-elementwise family, op 0..RT_EW_OPS-1 (Mul Add Sub Div): the kernel-name suffix and the C operator. Shared with rt_kernel_census.

@ rt_ew_name i op → s

@ rt_ew_sym i op → s

@ rt_sigmoid * EngineImpl e ONode n → v

@ rt_reshape * EngineImpl e ONode n → v

Reshape: pure reinterpret (data is contiguous) → alias the input buffer under the output name with the new shape. Shape comes from the INT64 initializer input[1]; a -1 entry is inferred, 0 copies the input dim.

@ rt_slice * EngineImpl e ONode n → v

Resize (nearest, integer upscale). Scales come from the FLOAT init input[2] = [1,1,sh,sw]. Slice along ONE axis, unit step, bounds from int64 initializers (inputs: data, starts, ends, axes[, steps]) — the shape torch 2.12 / onnxsim emit for channel splits (older exports used Split, which has its own handler). Negative axes normalise; ends clamp to the dim.

@ rt_resize * EngineImpl e ONode n → v

@ rt_transpose * EngineImpl e ONode n → v

General transpose (≤6-D): input dims + perm straight to gkd_perm (which pads to its 6-D kernel with trailing 1s / identity).

@ rt_softmax * EngineImpl e ONode n → v

Softmax over axis, viewing the tensor as (outer, axis, inner).

@ rt_concat * EngineImpl e ONode n → v

Concat along axis, viewing each input as (outer, axis_i, inner).

@ rt_split * EngineImpl e ONode n → v

Split along axis into contiguous slices — alias each output onto the input buffer at its byte offset (no copy). Sizes from the INT64 init input[1] when present, else num_outputs equal parts.

@ rt_matmul * EngineImpl e ONode n → v

MatMul. Two cases: A[...,M,K] @ B[K,N] (2-D B) collapses leading dims into M and uses Gemm; A[...,M,K] @ B[...,K,N] (matching batch dims, the attention case) uses a batched matmul.

@ rt_layernorm * EngineImpl e ONode n → v

LayerNormalization over the last axis (scale = in1, bias = in2).

@ rt_erf * EngineImpl e ONode n → v

@ rt_gathernd * EngineImpl e ONode n → v

GatherND — specialised for the CLIP EOS read-out: data [B,L,D] and an index built from (arange, argmax(tokens)) selecting each row's EOS token. Rather than materialise the int64 index chain, gather directly from the graph's token input: out[b,:] = data[b, argmax(tokens[b]), :].

@ rt_argmax * EngineImpl e ONode n → v

Gather along axis. Two index sources: a device tensor (e.g. the token matrix → embedding lookup) gathers nidx rows; a host scalar initializer (e.g. the QKV split index) selects one slice and drops the axis. ArgMax along the last axis. The output is int64 (allocated at 8 bytes/elem — rt_alloc_out's 4-byte default would let the kernel write past the buffer). Dispatches to the int64 kernel when the input is the graph's raw token input (the only 8-byte tensor in play; everything else on the value map is f32). Used by the CLIP text encoder's EOT read-out (ArgMax over token ids → Gather), which onnxsim leaves as a plain ArgMax+Gather pair — the eos_gather fast path only matches the old GatherND formulation.

@ rt_gather * EngineImpl e ONode n → v

@ rt_einsum * EngineImpl e ONode n → v

Einsum "bchw,bkc->bkhw": region[1,C,H,W] · text[1,K,C] -> [1,K,H,W]. = Gemm(text[K,C], region[C,HW]) -> [K,HW].

@ rt_reducel2 * EngineImpl e ONode n → v

@ rt_clip * EngineImpl e ONode n → v

@ rt_expand * EngineImpl e ONode n → v

Expand the last axis (broadcast a (...,1) tensor to the target shape).

@ rt_unsqueeze * EngineImpl e ONode n → v

Unsqueeze: insert size-1 axes — pure reshape (alias). New shape from the input's shape with 1s inserted at the axes positions.

@ rt_run_shaped Engine e__h OGraph g * u input_host sink ( Vec i ) shape → RTensor

Run the graph on a host input buffer (raw f32). shape is the input tensor shape (e.g. [1,3,416,416]), consumed (the value map keeps it). Returns the output device tensor.

@ rt_run_tokens Engine e__h OGraph g * u tokhost i nrow i ncol → RTensor

Token run: the single input is an INT64 token matrix [nrow, ncol] already laid out in tokhost (8-byte LE). Uploaded as-is for the embedding Gather.

@ rt_run_two Engine e__h OGraph g s n1 * u h1 sink ( Vec i ) s1 s n2 * u h2 sink ( Vec i ) s2 → RTensor

Two-input run (e.g. image + text embeddings for a promptable model).

@ _rt_run_nodes * EngineImpl e OGraph g → RTensor

@ rt_run Engine e__h OGraph g * u input_host i in_rows i in_cols → RTensor

Convenience for a 2-D (dense) input.

@ rt_download Engine e__h RTensor t → GpuHost

Download a device tensor into a fresh host f32 buffer — a GpuHost, released with its last owner (gpu_host_ptr / gpu_host_get_f32 read it).

@ rt_output1 Engine e__h → RTensor

The model's SECOND output (segmentation proto) after a run — valid until the next rt_reset. nelem 0 if the model has no second output. The value map still holds it because reset only happens at the start of a run.

@ rt_close sink Engine e → v

Let go of e now rather than at the end of its owner's scope. The engine's last owner releases its device blocks, its value map and the device (the kit).

@ rt_kernel_census GpuKit kit → i