grad.nu

grad — reverse-mode automatic differentiation over tensor.

Define-by-run: every g_* op computes its value EAGERLY through the tensor package and appends one node to a flat TAPE. backward(loss) walks the tape once in reverse — node inputs always have smaller indices, so reverse index order IS topological order — accumulating dL/dx into a gradient slot per node. Fan-out (one value used twice) sums naturally.

Ownership: the tape is an ARENA behind a handle (a GTape, an rcbox — every copy is the same tape). It owns every value tensor and every gradient tensor; grad_param/grad_const COPY the caller's tensor in, gvar_value/grad_of hand out BORROWS (aliased Tensor views — never free them, invalid once the tape is gone or reset past the node), and the tape's last owner releases everything (tape_free is an optional early release). No per-node ownership, no per-node closures — a ( Vec GNode ) of {op, a, b, scalar} plus two parallel pointer vecs. This keeps the hot loop free of NURL's closure-capture hazards and gives a deterministic reverse order for free.

The minibatch pattern: register parameters once, take tape_mark, then per batch build the episode's graph, backward, step the optimizer, and tape_reset_to(mark) — parameters (and their Adam state, which lives in the optimizer) survive, intermediates are freed, the arena keeps its capacity.

M1 scope (CPU): elementwise ops + all-axes reductions. Binary ops require EQUAL shapes (broadcast backward is M2, with its own oracle); scalar variants g_adds/g_muls cover the mixed case. A shape mismatch POISONS the tape (tape_ok → F, ops keep returning without touching memory) rather than half-computing.

Verification: tests/grad_test.nu checks every backward rule two independent ways — central finite differences and hand-derived analytic identities (e.g. d/dx Σx² = 2x, exact to the bit).

API

@ gop_param → i

── op codes ──────────────────────────────────────────────────────────

@ gop_const → i

@ gop_add → i

@ gop_sub → i

@ gop_mul → i

@ gop_div → i

@ gop_neg → i

@ gop_adds → i

@ gop_muls → i

@ gop_relu → i

@ gop_sigmoid → i

@ gop_tanh → i

@ gop_exp → i

@ gop_log → i

@ gop_sqrt → i

@ gop_sum → i

@ gop_mean → i

@ gop_matmul → i

@ gop_bmm → i

@ gop_transpose → i

@ gop_reshape → i

@ gop_softmax → i

@ gop_slice → i

@ gop_concat → i

: GNode

: GNode {
    i op
    i a
    i b
    f s
}

One tape node: the op, its input node ids (-1 = none), a scalar operand.

: GVar

: GVar {
    i id
}

A variable on the tape — an opaque node handle.

: GTapeImpl

: GTapeImpl {
    i ok  // 1 healthy · 0 poisoned (shape mismatch etc.)
    ( Vec GNode ) nodes
    ( Vec s ) vals  // *Tensor per node (tape-owned)
    ( Vec s ) grads  // *Tensor per node, 0 until backward touches it
    ( Vec s ) aux  // *GAux per node, 0 for ops that need none (slice starts)
}

: GTape { s ctl }

A GTape is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same state, and the last owner releases it.

@ GTape_share GTape h → GTape

@ GTape_drop sink GTape h → v

@ _GTape_ptr GTape h → *GTapeImpl

The state, for this package's own code.

% Drop GTapeImpl

% Drop GTapeImpl {
    @ drop GTapeImpl x → v {
        : i n ( vec_len [s] . x vals )
        : ~ i k 0
        ~ < k n {
            ( _g_tfree ?? ( vec_get [s] . x vals k ) { T y → y F → # s 0 } )
            ( _g_tfree ?? ( vec_get [s] . x grads k ) { T y → y F → # s 0 } )
            ( _g_auxfree ?? ( vec_get [s] . x aux k ) { T y → y F → # s 0 } )
            = k + k 1
        }
    }
}

The arena's raw cells — the Tensor value / gradient per node and the GAux blocks — are the tape's own heap memory: its last owner releases them here, and the node / pointer vecs after (drop glue).

: GAux

: GAux {
    ( Vec i ) starts
}

Per-node auxiliary payload (only g_slice uses one today).

@ tape_new → GTape

@ _g_tfree s pp → v

@ _g_auxfree s pp → v

@ tape_free sink GTape tp → v

Let go of tp now rather than at the end of its owner's scope. The tape's last owner releases every value and gradient tensor and every aux block it holds (the Drop below).

@ tape_ok GTape tp__h → b

@ tape_drop_consts GTape tp__h → v

Free the host VALUE tensors of every const (gop_const) node — the frozen base weights and the input consts. After a device capture (gput_capture) has uploaded them, nothing on the host reads them again: device replay reads device buffers, param sync-back touches only gop_param nodes, and gput_set_input uploads freshly-recomputed rows straight to the device. So this reclaims the single largest host allocation in a finetune (the base weights, held f64) with no effect on device training.

ONLY call it after a successful device capture: the CPU tape can no longer forward/backward once its const values are gone. Freeing is null-safe — tape_free / tape_reset_to re-read through the same guarded idiom.

@ tape_len GTape tp__h → i

@ tape_mark GTape tp__h → i

Watermark for tape_reset_to: everything appended after mark is dropped.

@ tape_reset_to GTape tp__h i mark → v

Drop every node at id >= mark (freeing its tensors) and ZERO the gradients of what remains — a fresh episode over the surviving parameters. The vecs keep their capacity, so a minibatch loop does not re-malloc the arena.

@ _g_heap Tensor t → s

Move a Tensor VALUE into a fresh heap cell the tape owns. The value's Vec handles alias into the cell — the caller must NOT free the value after.

@ _g_val * GTapeImpl tp i id → s

@ _g_grad_ptr * GTapeImpl tp i id → s

@ _g_poison * GTapeImpl tp s why → GVar

@ grad_poison GTape tp__h s why → GVar

Poison the tape from outside (a builder that meets a shape it cannot express): reports why once and hands back the invalid GVar.

@ _g_push * GTapeImpl tp i op i a i b f sc Tensor val → GVar

Append a node whose value is val (ownership moves to the tape).

@ _g_same_shape s pa s pb → b

Equal-shape check for M1 binops.

@ grad_param GTape tp__h Tensor w → GVar

Register a PARAMETER (requires-grad leaf). The tensor is COPIED in; the live, optimizer-updated copy is the tape's (read it via gvar_value).

@ grad_const GTape tp__h Tensor c → GVar

Register a CONSTANT (no gradient flows into it).

@ grad_const_lazy GTape tp__h ( Vec i ) shape i dtype → GVar

A constant declared by SHAPE, with its values supplied later — after the device capture, through gput_set_input, one tensor at a time.

Why: a frozen base weight is read exactly once, when the capture uploads it. Building the graph eagerly still makes every one of them resident in host RAM as f64 at the same moment, so the host peak is the whole model however little of it the training step actually needs. Declared this way the graph costs a shape, and the caller streams the values in from wherever they live.

The CPU tape CANNOT evaluate through a lazy const — it has no values. Use it only on a graph that will be captured and run on the device.

@ gvar_value GTape tp__h GVar v → Tensor

The node's value as a BORROWED Tensor view (do not free; invalid after tape_free / a reset past this node).

@ grad_of GTape tp__h GVar v → Tensor

The accumulated gradient as a BORROWED Tensor view. Allocates a zero gradient on first touch so the borrow is always valid.

@ g_scalar GTape tp__h GVar v → f

First element of the node's value — the scalar-loss readout.

@ _g_view s pp → Tensor

A borrowed Tensor view of a tape value (alias — never free).

@ _g_binop * GTapeImpl tp GVar a GVar b i op → GVar

@ g_add GTape tp__h GVar a GVar b → GVar

@ g_sub GTape tp__h GVar a GVar b → GVar

@ g_mul GTape tp__h GVar a GVar b → GVar

@ g_div GTape tp__h GVar a GVar b → GVar

@ _g_unary * GTapeImpl tp GVar a i op f sc → GVar

@ g_neg GTape tp__h GVar a → GVar

@ g_adds GTape tp__h GVar a f sc → GVar

@ g_muls GTape tp__h GVar a f sc → GVar

@ g_relu GTape tp__h GVar a → GVar

@ g_sigmoid GTape tp__h GVar a → GVar

@ g_tanh GTape tp__h GVar a → GVar

@ g_exp GTape tp__h GVar a → GVar

@ g_log GTape tp__h GVar a → GVar

@ g_sqrt GTape tp__h GVar a → GVar

@ _g_reduce * GTapeImpl tp GVar a i op → GVar

@ g_sum GTape tp__h GVar a → GVar

@ g_mean GTape tp__h GVar a → GVar

@ g_mse GTape tp__h GVar y GVar t → GVar

Mean squared error as a composite: mean((y − t)²). Exercises fan-out-free chaining; the FD harness covers it end to end.

@ g_matmul GTape tp__h GVar a GVar b → GVar

@ g_bmm GTape tp__h GVar a GVar b → GVar

@ g_transpose GTape tp__h GVar a → GVar

@ g_reshape GTape tp__h GVar a ( Vec i ) shape → GVar

shape is borrowed (copied here; tensor_reshape adopts the copy).

@ g_softmax GTape tp__h GVar a i axis → GVar

@ g_slice GTape tp__h GVar a ( Vec i ) starts ( Vec i ) stops → GVar

starts/stops are borrowed; starts is retained (aux) for the backward scatter.

@ g_concat GTape tp__h GVar a GVar b i axis → GVar

@ _g_acc_reduce * GTapeImpl tp i dst Tensor c f sgn → v

Accumulate an out-shaped contribution c into node dst's gradient, summing over broadcast axes (numpy alignment: trailing dims line up, a size-1 or missing dst dim receives the sum over that out axis). sgn is +1/-1 (sub/div's second input negates). Row-major walk of c — one documented, deterministic accumulation order.

@ _g_ensure_grad * GTapeImpl tp i id → v

Allocate node id's gradient as zeros (same shape as its value) if absent.

@ backward GTape tp__h GVar loss → b

dL/d(node) of loss seeds to ones; every earlier node receives the sum of its consumers' contributions. Returns F on a poisoned/invalid tape.

---

gpfuse.nu

grad/gpfuse.nu — MEGAKERNEL fusion: generate fused kernels from a captured program. The per-node replay's wall on small graphs is GPU-side per-kernel latency (~40 small dependent kernels per episode even under a CUDA graph); this module derives aegpu's fused shape FROM THE TAPE:

ROW-SPACE — a node is row-local when output row s depends only on row s of its activation inputs plus whole leaf operands: elementwise add/sub/mul/div ([B,c] with [B,c], [c] or [1]), every unary, and matmul act[B,k] · W[k,c] with W a leaf. A maximal consecutive run of row-local nodes becomes ONE generated kernel: block = row s, threads stride each node's columns in node order, __syncthreads() between nodes. Every node still writes ITS OWN buffer, so readbacks, unfused consumers and gput_set_input work unchanged.

Everything else (reductions, transpose, bmm, slice/concat, the scalar loss chain) keeps its per-node kernel — a graph that fuses nothing simply runs exactly as before. Execution order is preserved because segments are consecutive id ranges inside the normal walk.

BIT-EXACTNESS: each fused element is computed by one thread with the same expression, intrinsics and inner-loop order as the per-node kernels, so fused output is bit-identical — the test gates on ==.

The generated source is emitted in f64 with the __d*_rn discipline; an f32 program reuses the same substitution pass as the stock kernels (type, intrinsics, libm) — kernel names carry a per-plan id, so gpukit's name-keyed cache never collides across programs or dtypes.

ABI: one pointer-table argument. The plan uploads every node's val- buffer address into an i64 table once; a kernel reads const T v<k> = (const T)(unsigned long long)vt[<k>]; which keeps every generated kernel's signature identical: (const long long* vt, long long B)

API

: GpPlanImpl

: GpPlanImpl {
    b ok
    i kid  // unique plan id baked into kernel names
    i rows  // B — the fused chain's row count
    ( Vec i ) segs  // flattened (lo, hi) inclusive node-id ranges
    String src  // every generated kernel, one source (dtype-adjusted)
    ( Vec String ) knames  // kernel name per segment
    ( Vec String ) bnames  // bwd row-kernel per segment (empty = per-node)
    ( Vec String ) pnames  // bwd param-kernel per segment (empty = none)
    ( Vec i ) pgrid  // param-kernel grid blocks per segment
    GkBuf vtab  // i64 table: node id → val-buffer device address
    GkBuf gtab  // i64 table: node id → grad-buffer device address (0 = none)
    GkBuf ftab  // i64 pairs (grad ptr, n) — the one-launch zero fill
    i fcnt  // pair count in ftab
    String fillk  // fill kernel name
    ( Vec i ) ssegs  // serial-space (lo, hi) ranges — the scalar tail
    ( Vec String ) snames  // serial fwd kernel per serial segment
    ( Vec String ) sbnames  // serial bwd kernel (empty = per-node)
}

: GpPlan { s ctl }

A GpPlan is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same state, and the last owner releases it.

@ GpPlan_share GpPlan h → GpPlan

@ GpPlan_drop sink GpPlan h → v

@ _GpPlan_ptr GpPlan h → *GpPlanImpl

The state, for this package's own code.

@ gpfuse_plan_ok GpPlan pl__h → b

Did the analysis find something to fuse?

@ gpfuse_plan_segs GpPlan pl__h → ( Vec i )

The fused segments as flattened (lo, hi) inclusive node-id ranges, and each segment's backward row-kernel name (empty = per-node). Borrowed: valid while the plan is.

@ gpfuse_plan_bnames GpPlan pl__h → ( Vec String )

: ~ i g_gpm_next 1

@ _gpf_rowlocal * GProgImpl pg i k i B → b

Is node k row-local over B rows? (see the header for the definition)

@ _gpf_t String o i k → v

@ _gpf_lit String o f v → v

(double)(<shortest round-trip literal>) — the f32 pass rewrites the cast, giving (float)(f64 literal) = the same value our uploads convert to.

@ _gpf_decl String o i k b mut → v

Declare const double* v<k> (or mutable for outputs) from the table.

@ _gpf_emit_seg * GProgImpl pg i lo i hi i B s kname String o → v

One fused-segment kernel over nodes [lo, hi].

@ _gpf_in String o * GProgImpl pg i inid i C → v

a-side element of a binop/unary at (s, c).

@ _gpf_inb String o * GProgImpl pg i inid i C → v

b-side element: matched shape, per-row vector, or scalar.

@ _gpf_expr * GProgImpl pg i k String o → v

The per-element expression for fused node k — the EXACT per-node kernel arithmetic (gp_ew_bc / gp_scal / gp_trans / gp_matmul), spelled inline.

@ gpfuse_free sink GpPlan pl → v

Let go of pl now rather than at the end of its owner's scope. A plan owns its kernel names and sources and its device tables, and releases them with its last owner.

@ gpfuse_plan GProg pg__h → GpPlan

Analyze the program, emit the fused-forward kernels, upload the pointer table. ok=F (with everything freed safe) when nothing fuses.

@ _gpfuse_fwd_launches * GProgImpl pg * GpPlanImpl pl → b

Launches only (no sync policy) — shared by the direct path and the CUDA-graph capture.

@ gpfuse_forward GProg pg__h GpPlan pl__h → b

Fused forward: segments as one kernel each, everything else per-node.

@ _gpf_g String o i k → v

@ _gpf_gdecl String o i k → v

double g<k> = (double)(unsigned long long)gt[<k>];

@ _gpf_seen ( Vec i ) seen i id → b

@ _gpf_gel String o i id i C → v

g<id>[s * C + c]

@ _gpf_same * GProgImpl pg i inid i B i C → b

Is ew input inid the same [B,c] shape as consumer nd?

@ _gpf_acc_open String o * GProgImpl pg i tid i C s sgn → v

accred wrapper: TARGET = __dadd_rn(TARGET, __dmul_rn(<sgn>, <src...>

@ _gpf_stage_open String o i W → v

One row-space backward stage loop for c < W around body pushes.

@ _gpf_stage_close String o → v

@ _gpf_emit_bwd_stages * GProgImpl pg i lo i hi i B String o ( Vec i ) vids ( Vec i ) gids → i

Row-space bwd kernel body for segment [lo,hi]; stage count returned.

@ _gpf_param_jobs * GProgImpl pg i lo i hi i B ( Vec i ) jobs → v

Param-space jobs of segment [lo,hi], descending consumer order, encoded as (target, kind, consumer) triples — kind 0 = matmul dW, 1 = ew b-side.

@ _gpf_bwd_ok * GProgImpl pg i lo i hi i B → b

Order safety: no gradient may collect BOTH row-space and param-space contributions inside one segment (the split would reorder them).

@ _gpf_emit_param * GProgImpl pg i lo i hi i B String body ( Vec i ) vids ( Vec i ) gids → i

The param-space kernel: one grid-stride stage per distinct target, jobs serial inside each thread — gp_bw_mm_b / gp_bw_accred element order. Returns the widest target n (grid sizing); 0 when there are no jobs.

@ _gpf_emit_bwd_seg * GProgImpl pg i lo i hi i B s kname String o → b

Assemble the row-space bwd kernel; F when the segment has no stages.

@ _gpf_emit_param_seg * GProgImpl pg i lo i hi i B s kname String o → i

Assemble the param-space kernel; returns the widest target n (0 = none).

@ _gpf_emit_fill s kname String o → v

The one-launch gradient zero fill over an (ptr, n) pair table.

@ _gpfuse_bwd_launches * GProgImpl pg * GpPlanImpl pl → b

Fused backward: one zero-fill launch, the seed, then the reverse walk with row-space + param-space kernels standing in for fused segments.

@ gpfuse_backward GProg pg__h GpPlan pl__h → b

@ gpfuse_worthwhile GProg pg__h → b

Fusion pays where kernel LAUNCH LATENCY dominates — the cuda backend. The cpu backend has no launch latency and simulates a block's threads as fibers, so a barrier-heavy fused kernel only adds swapcontext cost there; the per-node path is already the optimal cpu execution. The planner still builds (the bit gates run everywhere) — this is the production selector.

@ gpfuse_graph_capture_train GProg pg__h GpPlan pl__h GpOpt go__h → b

The megakernel episode as ONE CUDA graph: fused forward + fused backward + the optimizer, captured once. Per episode the host does gput_set_input + gpopt_prepare + gput_episode — the identical driver loop as the per-node graph, just with ~5 kernels inside instead of ~40.

@ _gpf_smax → i

@ _gpf_in_rowseg * GpPlanImpl pl i k → b

@ _gpf_serial_ok * GProgImpl pg i k → b

@ _gpf_sel String o i id s ix → v

v<id>[<ix>] / g<id>[<ix>]

@ _gpf_sgel String o i id s ix → v

@ _gpf_sinb String o * GProgImpl pg i inid s ix → v

b-side element with the scalar-[1] broadcast collapsed to [0].

@ _gpf_emit_ser_stages * GProgImpl pg i lo i hi String o ( Vec i ) vids → i

Serial forward stages; node values in program order, exact expressions.

@ _gpf_sexpr * GProgImpl pg i k String o → v

The serial per-element expression at flat index e (gp_ew_bc / gp_scal / gp_trans forms with the trivial layouts this class permits).

@ _gpf_emit_ser_bwd_stages * GProgImpl pg i lo i hi String o ( Vec i ) vids ( Vec i ) gids → i

Serial backward stages hi→lo — gp_bw_reduce / gp_bw_unary / accred element order, one thread, exact accumulate order per element.

@ _gpf_ser_bsrc * GProgImpl pg i k String o → v

mul/div/add/sub b-side source element (accred's scr chain, inline).

@ _gpf_ser_unary * GProgImpl pg i k String o → v

gp_bw_unary's accumulate expression with gi bound, serial index e.

@ _gpf_emit_ser_seg * GProgImpl pg i lo i hi s kname String o → b

Assemble the serial fwd kernel (T on success — always has stages).

@ _gpf_emit_ser_bwd_seg * GProgImpl pg i lo i hi s kname String o → b

Assemble the serial bwd kernel (F when no reach-1 stages).

: GpFuseImpl

: GpFuseImpl {
    b active  // T when a worthwhile plan drives the episode
    GpPlan pl
}

: GpFuse { s ctl }

A GpFuse is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same state, and the last owner releases it.

@ GpFuse_share GpFuse h → GpFuse

@ GpFuse_drop sink GpFuse h → v

@ _GpFuse_ptr GpFuse h → *GpFuseImpl

The state, for this package's own code.

@ gpfuse_open GProg pg__h GpOpt go__h → GpFuse

@ _gpfuse_nullplan → GpPlan

@ gpfuse_active GpFuse s__h → b

@ gpfuse_episode GpFuse s__h GProg pg__h GpOpt go__h → b

One training episode. With a captured graph this is a single launch; with a live plan it is the fused forward + backward + optimizer step; otherwise the per-node forward/backward/step. Bit-identical either way.

@ gpfuse_close sink GpFuse s → v

Let go of s now rather than at the end of its owner's scope (its plan goes with its last owner).

---

emitc.nu

grad/emitc.nu — emit a recorded SCALAR tape as CUDA-C source: the __device__ void grad(...) function swarm-mcp's compute_iterate runs on every worker. This closes the loop the keystone plan named: distributed backward passes are now DERIVED from the tape, not hand-written CUDA-C.

Scope: every node must be a single-element tensor (shape [1] — scalars; g_sum/g_mean of a scalar pass through). The per-example loss is built once with placeholder leaves; the emitter walks the tape and prints forward: double tN = <op mirror>; backward: double gN = 0.0; … gA += …; (reverse order, seed 1.0) epilogue: swarm_g_add(g, j, g<param_j>); Leaf binding: parameters map to p[j] in registration order; data leaves map to caller-supplied C expressions (v, (double)x, …). Requires-grad propagation mirrors grad.nu — const-only branches emit no backward code.

The C forms mirror grad.nu's scalar arithmetic EXPRESSION BY EXPRESSION (sigmoid = 1/(1+exp(0-x)), tanh via e2=exp(2x), div's gB via gN*tN/tB…), so a host build with -ffp-contract=off reproduces the tape BIT FOR BIT — tests/emitc_oracle.sh proves exactly that with gcc.

API

@ gemit_cuda_grad GTape tp__h GVar loss ( Vec i ) pids ( Vec i ) lids ( Vec s ) lexprs b has_v String out → b

Emit the whole device function. has_v picks the dataset-bearing signature compute_iterate uses. F on any non-scalar node / unknown op.

---

gput.nu

grad/gput.nu — the GPU replay engine: run a captured tape episode on the device, bit-exactly.

The CPU tape (grad.nu) stays the semantic reference. gput_capture takes ONE recorded episode — parameters, constants, every op node up to the loss — and mirrors it onto the device: a value buffer and a gradient buffer per node, plus the stride/offset metadata each kernel needs. After capture the graph is STATIC (the mark/reset training loop re-records the same node structure every episode anyway): per minibatch the caller uploads fresh input rows into their const slots (gput_set_input), replays forward and backward on the device, and steps the device optimizer — nothing but the loss scalar needs to come back to the host.

Bit-exactness discipline (the aegpu recipe, generalized per op):

• Every kernel reproduces the CPU implementation's floating-point

ROUNDING AND ORDER: accumulation chains are spelled with explicit __dadd_rn/__dsub_rn/__dmul_rn/__ddiv_rn (never fused — NVRTC's default fmad contraction would change results), inner products and reductions run serially in the CPU's documented index order inside one thread, and broadcast-reduce accumulation walks each slot's contributions in the same row-major subsequence order as grad.nu's gacc_reduce.

• Optimizer scalars that need libm pow (Adam's lr_t) are computed on the

HOST with the same expressions opt.nu uses and shipped to the kernel.

• exp/log-based unaries (sigmoid, tanh, exp, log, softmax) mirror the

CPU FORMULAS exactly, but the device's exp()/log() are CUDA libm: on the gpu package's CPU backend they are glibc and the results are bit-identical; on real CUDA hardware they can differ by a few ulp. Everything else — relu, +, −, ×, ÷, sqrt, matmul/bmm, sum/mean, data movement, SGD/Adam — is bit-exact on BOTH backends. An MLP/AE workload (relu + linear + mse + Adam) is therefore bit-exact end-to-end everywhere; tests/gput_parity_test.nu pins both tiers.

Coverage: every grad.nu op. Restrictions (capture poisons with a message): TE_F64 tapes only, and g_bmm only with both operands carrying the full batch (no batch broadcast).

Ownership: GProg owns every device buffer; one gput_free. The tape is borrowed during capture and NOT needed afterwards (except as the target of gput_param_sync_host, which writes trained parameter values back into the tape's own tensors so gvar_value keeps working).

API

@ _str_replace_all s hay s needle s rep → String

Every simple textual replacement of needle with rep in hay. needle must be non-empty; used only on the fixed kernel source above.

: GpNode

: GpNode {
    i op
    i a
    i b
    f s
    i n  // value element count
    i reach  // 1 = loss ancestor (backward touches it)
    i rows  // shape[0] when the value is 2-D, else 0 (fusion analysis)
    i cols  // shape[1] when 2-D; the length when 1-D; else 0
    ( Vec i ) dims
    GkBuf val
    GkBuf grad
    GkBuf scr  // out-shaped scratch (mul/div backward)
    GkBuf meta  // i64 stride/offset blocks
}

One mirrored node. dims caches per-op host-side launch geometry: binop broadcast: [offB1 offB2 offB3 offBA offBB tfreeA tfreeB] matmul [M K N] · bmm [B M K N] · transpose [M N] softmax/concat [outer axlen inner (da)]

: GProgImpl

: GProgImpl {
    b ok
    GpuKit kit
    ( Vec GpNode ) nodes
    i loss
    GpuGraph gexec  // the CUDA graph of one fwd+bwd episode (none until captured)
    i dtype  // 0 GK_F64 (default) · 1 GK_F32 (device replay in float32)
}

: GProg { s ctl }

A GProg is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same state, and the last owner releases it.

@ GProg_share GProg h → GProg

@ GProg_drop sink GProg h → v

@ _GProg_ptr GProg h → *GProgImpl

The state, for this package's own code.

: ~ i g_gp_src_f32 0

The f32 source is constant but building it runs 10 full-source string passes, so cache it — otherwise gpsrc (called on EVERY kernel launch, though only the NAME feeds gpukit's cache) would regenerate ~15 KB of string ops thousands of times per training and dominate the run.

: ~ i g_gp_src_f32_ctl 0 // the String itself, kept for the program's lifetime

: ~ i g_gp_src_mixed 0

: ~ i g_gp_src_mixed_ctl 0 // the String itself, kept for the program's lifetime

@ _gp_src * GProgImpl pg → s

Kernel source for a program's dtype.

: GpOptImpl

: GpOptImpl {
    b ok
    i kind  // 0 sgd · 1 adam
    f lr
    f clip
    i t
    ( Vec i ) ids
    ( Vec f ) alphas
    ( Vec GkBuf ) m
    ( Vec GkBuf ) v
    GkBuf ptab  // [dptr len]* per param (built lazily for the clip norm)
    GkBuf nrm  // 1-elem norm result
    GkBuf ctl  // [lrt, cs] the update kernels read (graph-capturable)
}

Device optimizer over a captured program's parameters — opt.nu mirrored: per-param L2, global-norm clip, the same runtime 1−β Adam arithmetic, the step counter behind the heap pointer.

: GpOpt { s ctl }

A GpOpt is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same state, and the last owner releases it.

@ GpOpt_share GpOpt h → GpOpt

@ GpOpt_drop sink GpOpt h → v

@ _GpOpt_ptr GpOpt h → *GpOptImpl

The state, for this package's own code.

@ _gp_nobuf → GkBuf

@ _gp_push_vec ( Vec i ) dst ( Vec i ) src → v

Push nd entries of v (aligned to nd out dims, missing lead = filler).

@ _gp_accred_block ( Vec i ) meta ( Vec i ) oshape ( Vec i ) ost ( Vec i ) eff * u tfree → v

The acc-reduce block for one input: [nd ost dec fdim fdec], plus tfree. dec = row-major strides of the input's dims aligned into out dims (1 on broadcast dims); fdim/fdec describe the odometer over the broadcast dims.

@ _gp_ew_block ( Vec i ) meta ( Vec i ) ost ( Vec i ) xe ( Vec i ) ye → v

An ew block [nd ost xe ye].

@ _gp_fail * GProgImpl pg s why → v

@ _gp_lazy * Tensor t → b

Upload a CPU tensor's data into a fresh device buffer of the given dtype. GK_F64 is a straight memcpy; GK_F32 converts on the device (see above). A LAZY const (grad_const_lazy) carries a shape and no values: allocate its buffer from the shape and upload nothing — the caller fills it after the capture via gput_set_input.

@ _gp_upload_tensor GpuKit kit * Tensor t i edt → GkBuf

@ gput_capture GpuKit kit GTape tp__h GVar loss → GProg

Mirror one recorded episode ([0, tape_len)) onto the device. loss is the node backward seeds from. Check gput_ok before use.

@ gput_capture_dt GpuKit kit GTape tp__h GVar loss i dtype → GProg

As gput_capture, but the DEVICE replay runs in the given element dtype (0 GK_F64 · 1 GK_F32). The CPU tape stays f64: an f32 program halves device memory and runs f32 ALUs, at float32 precision (the parity test bounds the gap ~1e-5 rel; it is NOT bit-equal to the tape like f64 is).

@ gput_ok GProg pg__h → b

@ gput_free sink GProg pg → v

Let go of pg now rather than at the end of its owner's scope. The program's last owner releases every node's device buffers (back to the kit's pool), the captured graph and its hold on the kit.

@ _gp_node * GProgImpl pg i id → GpNode

@ _gp_dl GpuKit kit GkBuf b ( Vec f ) dst → b

dtype-aware transfers: f64 buffers take gk's straight-memcpy path, big f32 buffers the staged cast. The 4096-element floor keeps small vectors (norms, biases) off the kernel-launch overhead.

@ _gp_ul GpuKit kit GkBuf b ( Vec f ) src → b

@ gput_set_input GProg pg__h GVar v ( Vec f ) data → b

Fresh values for an input slot (a const's minibatch rows, or a param).

@ _gp_run * GProgImpl pg s name i grid i block ( Vec i ) args → b

@ _gp_fill_buf * GProgImpl pg GkBuf b f v → b

@ _gp_fwd_node * GProgImpl pg i k → b

One forward launch for node k. Leaves are data — nothing to do.

@ gput_forward GProg pg__h → b

Recompute every non-leaf value on the device, in tape order.

@ _gp_accred * GProgImpl pg GpNode dst GpNode nd GkBuf src i moff i tfree f sgn → b

The out-shaped contribution src accumulated into input dst's gradient over the broadcast axes (block at moff, odometer size tfree).

@ _gp_scr_ew * GProgImpl pg GpNode nd GkBuf x GkBuf y i moff i op → b

g (out-shaped) ⊙ other-input → scr, via the ew block at moff.

@ _gp_bwd_node * GProgImpl pg i k → b

@ gput_backward GProg pg__h → b

@ gput_graph_capture GProg pg__h → b

── graph fusion: one episode (forward + backward) as ONE launch ──── CUDA only; every kernel, argument value and launch order is IDENTICAL to the per-node path (bit-identical results) — the graph merely removes ~2N launch round-trips per episode. Inputs still refresh normally with gput_set_input (plain HtoD; the recorded kernels read the buffers). Returns F where graphs are unavailable (CPU backend) — callers keep using gput_forward/gput_backward.

@ gput_graph_capture_train GProg pg__h GpOpt go__h → b

Capture a whole TRAINING episode — forward + backward + the optimizer update — as one graph. Per episode the host then does: refresh inputs (gput_set_input), gpopt_prepare (one 16-byte ctl upload), and ONE gput_episode launch. The optimizer's per-step scalars ride the ctl buffer, so the captured launches never change.

@ gput_episode GProg pg__h → b

Run one fused episode (falls back to the per-node path without a graph).

@ gput_value GProg pg__h GVar v ( Vec f ) out → b

Download a node's value into out (resized by the caller to node size).

@ gput_grad GProg pg__h GVar v ( Vec f ) out → b

@ gput_loss GProg pg__h → f

The loss scalar (node value element 0).

@ gput_param_sync_host GProg pg__h GTape tp__h → b

Write every parameter node's device value back into the CPU tape's own tensor (in place), so gvar_value reflects the trained weights.

@ gpopt_new i kind f lr → GpOpt

@ gpopt_sgd_new f lr → GpOpt

@ gpopt_adam_new f lr → GpOpt

@ gpopt_free sink GpOpt o → v

Let go of o now rather than at the end of its owner's scope (its moment and table buffers go with its last owner).

@ gpopt_set_clip GpOpt o__h f maxn → v

@ gpopt_t GpOpt o__h → i

@ gpopt_set_t GpOpt o__h i t2 → v

@ gpopt_count GpOpt o__h → i

@ gpopt_m_download GpOpt o__h GProg pg__h i pi ( Vec f ) out → b

@ gpopt_v_download GpOpt o__h GProg pg__h i pi ( Vec f ) out → b

@ gpopt_m_upload GpOpt o__h GProg pg__h i pi ( Vec f ) src → b

@ gpopt_v_upload GpOpt o__h GProg pg__h i pi ( Vec f ) src → b

@ gpopt_add GpOpt o__h GProg pg__h GVar p f alpha → v

Register one parameter (its Adam moments start at zero, on the device).

@ _gpopt_ensure * GpOptImpl o * GProgImpl pg → b

Ensure the ctl buffer and (with clipping) the [dptr len] table exist.

@ gpopt_prepare GpOpt o__h GProg pg__h → b

The HOST half of a step: advance t, compute Adam's lr_t (host pow, like opt.nu), and upload [lrt, 1.0] into ctl. The device half then overwrites ctl[1] with the clip scale when clipping is on. Graph-safe: per episode this is ONE 16-byte upload.

@ _gpopt_launches * GpOptImpl o * GProgImpl pg → b

The DEVICE half: (clip-scale kernel when clipping) + one gp_opt launch per backward-active parameter. Pure launches — capturable into a graph.

@ gpopt_step GpOpt o__h GProg pg__h → b

One update step from the gradients on the device — opt.nu's semantics: untouched params (not loss ancestors) are SKIPPED, the clip norm covers every registered ancestor's gradient in registration order, Adam's lr_t comes from host pow exactly like the CPU path.

---

opt.nu

grad/opt.nu — optimizers over tape parameters: SGD and Adam, per-parameter L2 (weight decay), optional global-norm gradient clipping.

An optimizer is built empty (opt_adam_new lr / opt_sgd_new lr) and parameters are added one by one with their OWN L2 coefficient (opt_add o tp p alpha) — weights typically carry alpha, biases 0, which is the param-groups need in miniature. opt_step o tp reads each parameter's accumulated gradient from the tape and updates the parameter value IN PLACE on the tape (the tape owns the live copy; read it back with gvar_value).

Semantics: g = grad + alpha·w (L2 applied to the raw gradient; the loss is expected to be mean-scaled — no hidden batch divide here) SGD : w -= lr·g Adam : m = β1·m + (1−β1)·g ; v = β2·v + (1−β2)·g² w -= lr·√(1−β2ᵗ)/(1−β1ᵗ) · m/(√v + ε) β1 = 0.9, β2 = 0.999, ε = 1e-8 — sklearn/PyTorch defaults.

The step counter lives behind the *OptImpl heap pointer, so it ADVANCES — the frozen-Adam-t bug class (a scalar counter field on a by-value struct never persists in NURL) cannot recur, and tests/opt pin the trajectory bit-for- bit against a hand-computed two-step Adam.

Clipping (opt_set_clip o maxn): one global L2 norm over every added parameter's gradient; if it exceeds maxn every gradient is scaled by maxn/norm before the update — the standard clip-by-global-norm.

A parameter whose gradient was never touched by backward() is SKIPPED (PyTorch's None-grad rule), not decayed.

API

: OptImpl

: OptImpl {
    i kind  // 0 sgd · 1 adam
    f lr
    f clip  // 0 = off
    i t  // Adam step count (advances — shared state)
    i total  // total moment elements
    ( Vec i ) ids  // param node ids
    ( Vec f ) alphas  // per-param L2
    ( Vec i ) offs  // each param's slab offset into m / v
    ( Vec i ) lens  // each param's element count
    ( Vec f ) m  // first moments (adam)
    ( Vec f ) v  // second moments (adam)
}

: Opt { s ctl }

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

@ Opt_share Opt h → Opt

@ Opt_drop sink Opt h → v

@ _Opt_ptr Opt h → *OptImpl

The state, for this package's own code.

@ _opt_new i kind f lr → Opt

@ opt_sgd_new f lr → Opt

@ opt_adam_new f lr → Opt

@ opt_free sink Opt o → v

Let go of o now rather than at the end of its owner's scope (its moments and parameter lists go with its last owner).

@ opt_t Opt o__h → i

Adam's step count so far (0 for SGD and before the first step).

@ opt_set_clip Opt o__h f maxn → v

@ opt_add Opt o__h GTape tp__h GVar p f alpha → v

Register one tape parameter with its L2 coefficient (0 for biases).

@ _opt_gnorm * OptImpl o * GTapeImpl tp → f

The global L2 norm of every registered parameter's gradient (0-grads skip).

@ opt_step Opt o__h GTape tp__h → v

One update step from the gradients currently on the tape.