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

@ load_f32 s path → ( Vec u )

Load a raw little-endian f32 file: its values as bytes (4 per element, a trailing partial one dropped), 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 on its state in a Locked — every copy is the same engine, one run at a time, and its last owner releases the value map and the device. Every device tensor owns its block (a gpukit GkBuf, handed back to the kit's pool by its last owner): rt_reset lets a run's blocks go early, and a tensor a caller still holds keeps its own. Nothing is freed by hand (rt_close is an optional early release). The graph a run is given is BORROWED for the run; the caller keeps it.

API

: RTensor { String name GkBuf buf ( Vec i ) shape i nelem b hosti }

A tensor of a run: name, device buffer, shape, element count. The buffer owns its block (or shares one with the tensor it aliases, for a Reshape or a Split), so an RTensor stays readable after the run that made it — rt_download of an old one reads what that run computed. A host-int tensor (hosti, see below) has no device buffer: its values are its shape.

: EngineImpl

: EngineImpl {
    GpuKit kit
    ( Vec RTensor ) vals
    String input_name
    String output_name
    String output1_name
    b ok
}

The engine's state: the kit (its device), the value map of the current run, and copies of the three graph names it needs after the run. The graph itself stays the CALLER's: a run borrows it, and the engine holds no part of it — what lets a caller run the same graph again, on this engine or another.

: Engine { ( Locked EngineImpl ) l GpuKit kit }

An Engine is a handle on its state in a Locked (stdlib/std/thread.nu): every copy is the same engine, a run takes the state for its length (another thread's run waits), and the last owner releases it. kit is another owner of the state's kit, for what reads it without a run.

@ _rt_none → RTensor

@ streq2 s a s b → b

@ ceil_div i a i b → i

@ 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 inout EngineImpl e OGraph g → v

Start a run of the caller's graph: the last run's tensors go, and the graph's names are copied (see EngineImpl). The graph stays the caller's.

@ rt_reset Engine h → v

Let the previous run's tensors go — their device blocks return to the kit's pool, unless a caller still holds one. Lets one Engine serve many forward passes (e.g. one text prompt per call) without holding the last pass's weights + activations.

@ rt_ok Engine h → b

@ rt_name Engine h → s

@ rt_put inout EngineImpl e s name GkBuf buf sink ( Vec i ) shape → v

Register a device tensor under name with an explicit shape vector. Its buffer is exactly the shape's elements: a view of the front of buf when that is longer (a Reshape to fewer elements), no buffer at all when it is shorter.

@ rt_find EngineImpl e s name → i

@ rt_at EngineImpl e i idx → RTensor

@ rt_own RTensor t → RTensor

A tensor of its own: the same device block (a counted share), its own name and shape — what a caller keeps after the engine goes back into its Locked, where the next run replaces the value map (rt_at is a view of it).

@ rt_load_inits inout EngineImpl e OGraph g → v

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

Allocate a fresh f32 device tensor with shape, register it under name, and hand back its buffer (another owner of the block).

@ rt_gemm inout EngineImpl e ONode n → v

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

@ rt_relu inout EngineImpl e ONode n → v

@ rt_conv inout EngineImpl e ONode n → v

@ rt_convtranspose inout 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 inout EngineImpl e ONode n → v

@ rt_batchnorm inout EngineImpl e ONode n → v

@ rt_leakyrelu inout EngineImpl e ONode n → v

@ rt_ndim RTensor t → i

@ rt_last RTensor t → i

@ rt_binop inout 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 inout EngineImpl e ONode n → v

@ rt_reshape inout EngineImpl e OGraph g 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 inout EngineImpl e OGraph g 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 inout EngineImpl e OGraph g ONode n → v

@ rt_transpose inout 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 inout EngineImpl e ONode n → v

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

@ rt_concat inout EngineImpl e ONode n → v

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

@ rt_split inout EngineImpl e OGraph g 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 inout 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 inout EngineImpl e ONode n → v

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

@ rt_erf inout EngineImpl e ONode n → v

@ rt_gathernd inout 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 inout 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 inout EngineImpl e OGraph g ONode n → v

@ rt_einsum inout 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 inout EngineImpl e ONode n → v

@ rt_clip inout EngineImpl e OGraph g ONode n → v

@ rt_expand inout EngineImpl e OGraph g ONode n → v

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

@ rt_unsqueeze inout EngineImpl e OGraph g 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 h OGraph g ( Slice u ) input sink ( Vec i ) shape → RTensor

Run the graph on host input bytes (f32, little-endian) shaped shape (e.g. [1,3,416,416]; consumed — the value map keeps it). Returns the output tensor; an empty one (nelem 0) when input is shorter than the shape says, or a dim is negative, and nothing ran.

@ rt_run_host Engine h OGraph g GpuHost input sink ( Vec i ) shape → RTensor

The same from a GpuHost (gpu_host_alloc, filled with gpu_host_set_f32).

@ rt_run_tokens Engine h OGraph g ( Slice u ) tok i nrow i ncol → RTensor

Token run: the single input is an INT64 token matrix [nrow, ncol] (8-byte LE), uploaded as-is for the embedding Gather; one too small runs nothing (an empty tensor).

@ rt_run_tokens_host Engine h OGraph g GpuHost tok i nrow i ncol → RTensor

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

Two-input run (e.g. image + text embeddings for a promptable model), each f32 bytes holding its shape's worth; either one too small runs nothing (an empty tensor).

@ rt_run_two_host Engine h OGraph g s n1 GpuHost h1 sink ( Vec i ) s1 s n2 GpuHost h2 sink ( Vec i ) s2 → RTensor

@ _rt_run_nodes inout EngineImpl e OGraph g → RTensor

@ rt_run Engine h OGraph g ( Slice u ) input i in_rows i in_cols → RTensor

Convenience for a 2-D (dense) input.

@ rt_download Engine h RTensor t → GpuHost

Download a tensor into a fresh host buffer — a GpuHost, released with its last owner (gpu_host_get_f32 / gpu_host_f32 read it). The tensor owns its device block, so one from an earlier run reads what that run computed; one with no buffer comes back empty. The bytes go straight into the pinned block (gk_dbuf_download_host): a GpuHost's memory is owned through an Arc, so writing it through a ( Slice u ) view of the payload is a change to frozen memory, which the borrow checker refuses.

@ rt_output1 Engine h → RTensor

The model's SECOND output (segmentation proto) after a run — empty (nelem 0) if the model has no second output.

@ 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

---

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

---

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.

: 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.

---

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 registers another owner of the input DTensor's

device block as the graph input: the block lives as long as the run needs it, whatever the caller does with the DTensor.

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

owned DTensor (released with its last owner). A copy, not a share, because a graph output may ALIAS another tensor's block at an offset (Reshape/Split), and a DTensor is the caller's to write.

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

API

@ rt_run_dtensor Engine 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.

@ dtensor_from_output GpuKit kit Engine e RTensor t → DTensor

Copy a run's output into an OWNED device tensor (device-to-device — see the ownership note above). The result lives independently of the engine and of the output.

---

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

@ onnx_static_c_from GpuKit kit String out ( Vec String ) names → i

Build kernels_static.c from a census kit's record into out; the emitted entry names go to names. Returns the number of problems (0 = the file is complete).

@ onnx_static_kernels_c String out ( Vec String ) names → i

The whole file, from a fresh census: kernels_static.c into out, the emitted entry names into names. Returns the number of problems — a census call gpukit rejected or a kernel the static backend cannot run — each reported on stderr; 0 = complete.