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 frees it — the caller still dtensor_frees it, but only AFTER the run's outputs are consumed or copied out).

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

owned DTensor. A copy (not an ownership transfer) because graph outputs may ALIAS an owned base allocation at an offset (Reshape/Split); freeing such a pointer would corrupt the heap, and the engine frees its own allocations at the next rt_reset anyway.

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

API

@ rt_run_dtensor * Engine e 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 / rt_close do not touch it; free with dtensor_free.

---

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
    i host  // *u as i64 (0 = 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. host == 0 for a graph value placeholder.

: 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

---

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 * u pcount → *u

Load a raw little-endian f32 file into a fresh host buffer; returns the buffer and writes the element count through pcount (an i64 cell).

@ 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).

API

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

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

: Engine { Gpu g * GpuKit kit ( Vec RTensor ) vals OGraph graph b ok ( Vec i ) owned b graph_set }

Engine.g BORROWS the kit's device handle (gk_close releases it).

@ 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 * Engine e OGraph g → v

Install the caller's graph on the engine. The placeholder OGraph that rt_open allocated is freed on the first install; caller graphs are BORROWED (the engine never frees them).

@ rt_own * Engine e i dptr → i

Record a device allocation so rt_reset can free it. Aliased tensors (Reshape/Split share or offset an existing buffer) are NOT recorded — only the real gpu_alloc base pointers, so reset frees each allocation exactly once with no double-free / free-of-non-base.

@ rt_reset * Engine e → v

Free 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 leaking the model's weights + activations each time.

@ rt_ok * Engine e → b

@ rt_name * Engine e → s

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

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

@ rt_find * Engine e s name → i

@ rt_at * Engine 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 * Engine e OGraph g → v

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

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

@ rt_gemm * Engine e ONode n → v

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

@ rt_relu * Engine e ONode n → v

@ rt_conv * Engine e ONode n → v

@ rt_convtranspose * Engine 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 * Engine e ONode n → v

@ rt_batchnorm * Engine e ONode n → v

@ rt_leakyrelu * Engine e ONode n → v

@ rt_ndim RTensor t → i

@ rt_last RTensor t → i

@ rt_binop * Engine 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.

@ rt_sigmoid * Engine e ONode n → v

@ rt_reshape * Engine 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 * Engine 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 * Engine e ONode n → v

@ rt_transpose * Engine 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 * Engine e ONode n → v

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

@ rt_concat * Engine e ONode n → v

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

@ rt_split * Engine 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 * Engine 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 * Engine e ONode n → v

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

@ rt_erf * Engine e ONode n → v

@ rt_gathernd * Engine 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 * Engine 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 * Engine e ONode n → v

@ rt_einsum * Engine 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 * Engine e ONode n → v

@ rt_clip * Engine e ONode n → v

@ rt_expand * Engine e ONode n → v

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

@ rt_unsqueeze * Engine 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 OGraph g * u input_host ( 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]). Returns the output device tensor.

@ rt_run_tokens * Engine e 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 OGraph g s n1 * u h1 ( Vec i ) s1 s n2 * u h2 ( Vec i ) s2 → RTensor

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

@ _rt_run_nodes * Engine e OGraph g → RTensor

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

Convenience for a 2-D (dense) input.

@ rt_download * Engine e RTensor t → *u

Download a device tensor into a fresh host f32 buffer (caller frees).

@ rt_output1 * Engine e → 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 * Engine e → v