pb.nu

packages/onnx/src/pb.nu — a minimal protobuf wire-format decoder.

Just enough proto3 to read an ONNX model: varints, the four wire types, and length-delimited regions (sub-messages, strings, bytes, packed repeated). Pure NURL over a ( Vec u ) byte buffer — no codegen, no schema; the ONNX message shapes live in model.nu, which drives this by field number.

A PbR is a cursor: the buffer plus the current read position. Decoders advance pos; message parsers loop while pos < end, where end comes from the enclosing length-delimited region (or the buffer length at top level). A sub-message is parsed by a fresh PbR over the same buffer (pb_sub) positioned at the payload — so nested ends never clash.

API

: PbR { ( Vec u ) buf i pos i len }

@ pb_wt_varint → i

Wire types (low 3 bits of a field tag).

@ pb_wt_i64 → i

@ pb_wt_len → i

@ pb_wt_i32 → i

@ pb_new ( Vec u ) buf → *PbR

New reader over a byte buffer (cursor at 0, end = buffer length).

@ pb_sub *PbR r i start i len → *PbR

A reader over the sub-region [start, start+len) of the SAME buffer.

@ pb_pos *PbR r → i

@ pb_set_pos *PbR r i p → v

@ pb_end *PbR r → i

@ pb_more *PbR r → b

@ pb_free *PbR r → v

@ pb_byte *PbR r i off → i

Byte at absolute offset (0 past end).

@ pb_varint *PbR r → i

Read a base-128 varint at the cursor; advance past it.

@ pb_tag *PbR r → i

Field tag → field number / wire type.

@ pb_field i tag → i

@ pb_wire i tag → i

@ pb_blob_start *PbR r → i

Length-delimited: read length, return payload start, advance cursor PAST it. Caller reads [start, start+len) (recover len with pb_sub or by math).

@ pb_submsg *PbR r → *PbR

Read a length-delimited field and return a bounded sub-reader over its payload, advancing the parent cursor past it. The caller loops the sub-reader with pb_more and frees it with pb_free. This is how embedded messages (and packed repeated fields) are descended into.

@ pb_i32 *PbR r → i

Fixed 32-bit little-endian at cursor; advance 4. (float bits / fixed32)

@ pb_string *PbR r → String

Read a length-delimited region as a String.

@ pb_skip *PbR r i wire → v

Skip a field of the given wire type (cursor at the value).

@ pb_read_f32_into *PbR r *u dst i n → v

Read n little-endian f32 values from the byte region at the cursor into the host buffer dst (raw u, 4-byte stride), writing each one's exact 32-bit pattern (no float round-trip). Advances the cursor by n4. Used for TensorProto.raw_data → host weights, ready to upload to the GPU.

@ pb_read_i64_into *PbR r *u dst i n → v

Read n little-endian int64 values from the byte region at the cursor into dst (8-byte stride). For TensorProto raw_data of an INT64 tensor (shape tensors, split sizes, anchor grids). Advances the cursor by n*8.

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

4-byte typed store (runtime.c accessor; runtime.o is always linked).

---

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 — ONNX schema parsing over pb.nu.

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 read straight into a host buffer ready to upload to the GPU.

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_parse ( Vec u ) bytes → OGraph

── ModelProto (top level) ────────────────────────────────────────

@ graph_find_init OGraph g s name → i

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

@ graph_free OGraph g → v

---

ops.nu

packages/onnx/src/ops.nu — ONNX operators as GPU kernels.

Each op is a CUDA-C kernel compiled once via the gpu package's NVRTC path and launched over the gpu.nu interface. Tensors live on the device as raw CUdeviceptr (i64); shapes are tracked by the executor (runtime.nu).

Op set: Gemm, Relu (dense MLP) + Conv, MaxPool, BatchNormalization, LeakyRelu, and broadcast Mul/Add (CNNs, NCHW) — enough to run a tiny-YOLO-class detector. All launches are 1-D: thread idx → output element; batch N is assumed 1.

API

: Kernels

: Kernels {
    GpuKernel gemm GpuKernel relu
    GpuKernel conv GpuKernel pool
    GpuKernel bn GpuKernel lrelu GpuKernel elt
    GpuKernel sig GpuKernel resize GpuKernel perm
    GpuKernel softm GpuKernel cpyax
    GpuKernel rl2 GpuKernel clip GpuKernel expand GpuKernel slice
    GpuKernel lnorm GpuKernel erf GpuKernel gather GpuKernel bmm GpuKernel amax GpuKernel amax64
    GpuKernel eos GpuKernel convt
    b ok
}

Compiled kernels, built once and reused across nodes.

@ ops_compile Gpu g → Kernels

@ ops_free Kernels ks → v

@ op_eos_gather Gpu g Kernels ks i datad i tokd i yd i B i L i D → i

@ op_layernorm Gpu g Kernels ks i xd i scd i bd i yd i outer i ax f eps → i

@ op_erf Gpu g Kernels ks i xd i yd i n → i

@ op_gather Gpu g Kernels ks i dd i idd i yd i outer i axis_in i inner i nidx → i

@ op_bmm Gpu g Kernels ks i ad i bd i yd i batch i M i N i K → i

@ op_argmax Gpu g Kernels ks i xd i yd i outer i ax → i

@ op_argmax_i64 Gpu g Kernels ks i xd i yd i outer i ax → i

@ op_slice_ax Gpu g Kernels ks i sd i dd i outer i sz i inner i src_ax i soff → i

@ op_reducel2 Gpu g Kernels ks i xd i yd i outer i ax → i

@ op_clip Gpu g Kernels ks i xd i yd i n f lo f hi → i

@ op_expand Gpu g Kernels ks i xd i yd i outer i rep → i

@ op_gemm Gpu g Kernels ks i adptr i bdptr i cdptr i ydptr i M i N i K f alpha f beta i transB → i

@ op_relu Gpu g Kernels ks i xdptr i ydptr i n → i

@ op_conv Gpu g Kernels ks i xd i wd i bd i yd i Cin i H i W i Cout i kh i kw i OH i OW i ph i pw i sh i sw i hasB → i

@ op_convtranspose Gpu g Kernels ks i xd i wd i bd i yd i Cin i H i W i Cout i kh i kw i OH i OW i ph i pw i sh i sw i hasB → i

@ op_maxpool Gpu g Kernels ks i xd i yd i C i H i W i kh i kw i OH i OW i sh i sw i ph i pw → i

@ op_batchnorm Gpu g Kernels ks i xd i scd i bd i md i vd i yd i C i HW f eps → i

@ op_leakyrelu Gpu g Kernels ks i xd i yd i n f alpha → i

@ op_eltwise Gpu g Kernels ks i xd i bd i yd i n i per i op i bmode → i

@ op_sigmoid Gpu g Kernels ks i xd i yd i n → i

@ op_resize Gpu g Kernels ks i xd i yd i C i H i W i OH i OW i sh i sw → i

@ op_perm6 Gpu g Kernels ks i xd i yd i d0 i d1 i d2 i d3 i d4 i d5 i p0 i p1 i p2 i p3 i p4 i p5 → i

@ op_softmax Gpu g Kernels ks i xd i yd i outer i ax i inner → i

@ op_copy_ax Gpu g Kernels ks i sd i dd i outer i src_ax i inner i dst_ax i off → v

Copy one source into a concat output's axis slot at off.

---

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 GPU kernel in ops.nu; the named output is downloaded at the end. A value map (name → device tensor) threads activations between nodes.

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 Kernels ks ( Vec RTensor ) vals OGraph graph b ok ( Vec i ) owned b graph_set }

@ 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 and compile the kernels.

@ 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

@ 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). Pad missing trailing dims to size 1 with an identity perm so the 6-D kernel handles any rank.

@ 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 * 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