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 Arc — 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.

: GTape

: GTape {
    i ok  // 1 healthy · 0 poisoned (shape mismatch etc.)
    ( Vec GNode ) nodes
    ( Vec Tensor ) vals  // the value per node
    ( Vec Tensor ) grads  // the gradient per node (an empty tensor until touched)
    ( Vec u ) ghas  // per node: 1 once its gradient exists
    ( Vec ( Vec i ) ) aux  // per node: slice starts (g_slice), empty for other ops
}

@ _g_empty → Tensor

@ _g_v GTape tp i id → Tensor

Node id's value / gradient, lent (empty when out of range or absent).

@ _g_gv GTape tp i id → Tensor

@ _g_has_grad GTape tp i id → b

@ tape_new → GTape

@ 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 (boxes in Vecs: dropped with them).

@ tape_ok GTape tp → b

@ tape_drop_consts inout GTape tp → 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 → i

@ tape_mark GTape tp → i

Watermark for tape_reset_to: everything appended after mark is dropped.

@ tape_reset_to inout GTape tp i mark → v

Drop every node at id >= mark (and 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_poison inout GTape tp s why → GVar

@ grad_poison inout GTape tp 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 inout GTape 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 Tensor a Tensor b → b

Equal-shape check for M1 binops.

@ grad_param inout GTape tp 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 inout GTape tp Tensor c → GVar

Register a CONSTANT (no gradient flows into it).

@ grad_const_lazy inout GTape tp ( 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 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 inout GTape tp 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 GVar v → f

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

@ _g_binop inout GTape tp GVar a GVar b i op → GVar

@ g_add inout GTape tp GVar a GVar b → GVar

@ g_sub inout GTape tp GVar a GVar b → GVar

@ g_mul inout GTape tp GVar a GVar b → GVar

@ g_div inout GTape tp GVar a GVar b → GVar

@ _g_unary inout GTape tp GVar a i op f sc → GVar

@ g_neg inout GTape tp GVar a → GVar

@ g_adds inout GTape tp GVar a f sc → GVar

@ g_muls inout GTape tp GVar a f sc → GVar

@ g_relu inout GTape tp GVar a → GVar

@ g_sigmoid inout GTape tp GVar a → GVar

@ g_tanh inout GTape tp GVar a → GVar

@ g_exp inout GTape tp GVar a → GVar

@ g_log inout GTape tp GVar a → GVar

@ g_sqrt inout GTape tp GVar a → GVar

@ _g_reduce inout GTape tp GVar a i op → GVar

@ g_sum inout GTape tp GVar a → GVar

@ g_mean inout GTape tp GVar a → GVar

@ g_mse inout GTape tp 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 inout GTape tp GVar a GVar b → GVar

@ g_bmm inout GTape tp GVar a GVar b → GVar

@ g_transpose inout GTape tp GVar a → GVar

@ g_reshape inout GTape tp GVar a ( Vec i ) shape → GVar

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

@ g_softmax inout GTape tp GVar a i axis → GVar

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

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

@ g_concat inout GTape tp GVar a GVar b i axis → GVar

@ _g_acc_reduce inout GTape 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_acc_into GTape tp i dst Tensor c f sgn → v

…into a gradient already there (written in place through its buffer: the tape itself does not change).

@ _g_ensure_grad inout GTape tp i id → v

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

@ backward inout GTape tp 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.

SAFETY: the passes below write plain CUDA C, reaching node k's value as v<k> and its gradient as g<k> through a table of addresses. That text is never launched: __gpf_kernels rewrites every kernel into gpukit's checked subset first — v12[i] becomes LD(v12, i), g7[i] = x; becomes ST(g7, i, x);, each v<k> / g<k> a binding of node k's buffer, the early return guards an enclosing if (a block with SYNC() may not leave early), every single-statement if / for / while body braced — and the launches are recorded (gk_record_dev) once and replayed. A kernel the checker refuses fails the plan; none can reach memory outside the buffers it is bound to.

API

: GpfKern

: GpfKern {
    String name
    String body  // the checked body
    ( Vec i ) vids  // v<k> bindings: node k's value
    ( Vec i ) gids  // g<k> bindings: node k's gradient
    b hasb  // takes the row count as `B`
}

: GpPlan

: GpPlan {
    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
    ( Vec GpfKern ) ks  // every generated kernel, as a checked body
    ( Vec GkLaunch ) frec  // fused forward launches, recorded on the first forward
    ( Vec GkLaunch ) brec  // fused backward launches, recorded on the first backward
    ( Vec String ) fills  // the gradient zero-fill kernels (one block per gradient)
    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)
}

@ gpfuse_plan_ok inout GpPlan pl → b

A GpPlan is a handle on its state in an Arc (stdlib/std/arc.nu): every copy is the same state, and the last owner releases it. The state, for this package's own code. Did the analysis find something to fuse?

@ gpfuse_plan_segs inout GpPlan pl → ( 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 inout GpPlan pl → ( Vec String )

: ~ i g_gpm_next 1

@ _gpf_rowlocal inout GProg 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 inout GProg pg i lo i hi i B s kname String o → v

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

@ _gpf_in String o inout GProg pg i inid i C → v

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

@ _gpf_inb String o inout GProg pg i inid i C → v

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

@ _gpf_expr inout GProg 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 inout GProg pg → 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 inout GProg pg inout GpPlan pl → b

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

@ gpfuse_forward inout GProg pg inout GpPlan pl → 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 inout GProg 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 inout GProg 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 inout GProg 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 inout GProg 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 inout GProg 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 inout GProg 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 inout GProg 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 inout GProg 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_fills inout GProg pg inout GpPlan pl → v

The gradient zero fill: one block per gradient buffer, in kernels of at most 64 gradients each (named fillk_<n>, listed in pl.fills).

@ _gpf_run inout GProg pg ( Vec GpfKern ) ks inout ( Vec GkLaunch ) rs inout i ri s kname i rows i nb i blk → b

Launch generated kernel kname — recorded into rs the first time (slot ri), replayed after it. rows is B.

@ _gpfuse_bwd_launches inout GProg pg inout GpPlan 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 inout GProg pg inout GpPlan pl → b

@ gpfuse_worthwhile inout GProg pg → 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 inout GProg pg inout GpPlan pl inout GpOpt go → 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 inout GpPlan pl i k → b

@ _gpf_serial_ok inout GProg 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 inout GProg pg i inid s ix → v

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

@ _gpf_emit_ser_stages inout GProg pg i lo i hi String o ( Vec i ) vids → i

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

@ _gpf_sexpr inout GProg 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 inout GProg 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 inout GProg pg i k String o → v

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

@ _gpf_ser_unary inout GProg pg i k String o → v

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

@ _gpf_emit_ser_seg inout GProg 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 inout GProg pg i lo i hi s kname String o → b

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

: GpFuse

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

@ gpfuse_open inout GProg pg inout GpOpt go → GpFuse

A GpFuse is a handle on its state in an Arc (stdlib/std/arc.nu): every copy is the same state, and the last owner releases it. The state, for this package's own code.

@ _gpfuse_nullplan → GpPlan

@ gpfuse_active inout GpFuse s → b

@ gpfuse_episode inout GpFuse s inout GProg pg inout GpOpt go → 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 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

@ _gp_body i dtype s name → String

A kernel's body for the program's dtype. f64 and mixed run it as written: mixed binds f32 buffers, and a load widens to double and a store narrows exactly as the old pointer-typed source did. Pure f32 swaps the element type, the round-to-nearest intrinsics and the libm calls — the same substitution the f32 kernel source always went through, so its arithmetic is unchanged.

@ _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)]

: GProg

: GProg {
    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)
    // The recorded launches (gk_record_dev), built on the first replay:
    // per node forward and backward, and the gradient zeroing + seed.
    ( Vec ( Vec GkLaunch ) ) frec
    ( Vec ( Vec GkLaunch ) ) brec
    ( Vec GkLaunch ) zrec
}

: GpOpt

: GpOpt {
    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 acc  // the clip norm's running sum of squares (1 element)
    GkBuf ctl  // [lrt, cs] the update kernels read (graph-capturable)
    ( Vec GkLaunch ) orec  // the step's launches, recorded on the first step
}

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.

@ _gp_nobuf → GkBuf

A GpOpt is a handle on its state in an Arc (stdlib/std/arc.nu): every copy is the same state, and the last owner releases it. The state, for this package's own code.

@ _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 inout i 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 inout GProg 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 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 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 inout GProg pg → 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 inout GProg 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 inout GProg pg GVar v ( Vec f ) data → b

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

@ _gp_go inout GProg pg inout ( Vec GkLaunch ) rs inout GkCall c s name i nb i blk → b

Bind and check kernel name over c's bindings, record the launch into rs and make it: nb blocks of blk threads. F (the program marked failed, with the reason) when the kernel is refused or fails.

@ _gp_fill_rec inout GProg pg inout ( Vec GkLaunch ) rs GkBuf b f v → b

Buffer b filled with v, recorded into rs.

@ _gp_fill_buf inout GProg pg GkBuf b f v → b

Buffer b filled with v (not recorded).

@ _gp_fwd_node inout GProg pg i k → b

One forward launch for node k. Leaves are data — nothing to do. The launches are recorded the first time and replayed after it: the same kernels over the same buffers every episode, bound and checked once.

@ gput_forward inout GProg pg → b

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

@ _gp_accred inout GProg pg inout ( Vec GkLaunch ) rs 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). The out-shaped contribution src accumulated into input dst's gradient over the broadcast axes (block at moff, odometer size tfree).

@ _gp_scr_ew inout GProg pg inout ( Vec GkLaunch ) rs 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 inout GProg pg i k → b

Node k's backward launches, recorded the first time and replayed after.

@ gput_backward inout GProg pg → b

@ gput_graph_capture inout GProg pg → 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 inout GProg pg inout GpOpt go → 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 inout GProg pg → b

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

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

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

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

@ gput_loss inout GProg pg → f

The loss scalar (node value element 0).

@ gput_param_sync_host inout GProg pg GTape tp → 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 inout GpOpt o f maxn → v

@ gpopt_t inout GpOpt o → i

@ gpopt_set_t inout GpOpt o i t2 → v

@ gpopt_count inout GpOpt o → i

@ gpopt_m_download inout GpOpt o inout GProg pg i pi ( Vec f ) out → b

@ gpopt_v_download inout GpOpt o inout GProg pg i pi ( Vec f ) out → b

@ gpopt_m_upload inout GpOpt o inout GProg pg i pi ( Vec f ) src → b

@ gpopt_v_upload inout GpOpt o inout GProg pg i pi ( Vec f ) src → b

@ gpopt_add inout GpOpt o inout GProg pg GVar p f alpha → v

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

@ _gpopt_ensure inout GpOpt o inout GProg pg → b

Ensure the ctl buffer and (with clipping) the norm accumulator exist. The accumulator keeps the sum's precision: f32 for a pure-f32 program, f64 for f64 and mixed (whose norm always summed in double).

@ gpopt_prepare inout GpOpt o inout GProg pg → 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 inout GpOpt o inout GProg pg → b

The DEVICE half: (the clip scale when clipping) + one opt launch per backward-active parameter. Pure launches — capturable into a graph — recorded on the first step and replayed after it. The clip norm sums every active parameter's squared gradient in registration order, one parameter per launch into the accumulator: the order the single-thread kernel over a table of raw gradient addresses summed in.

@ gpopt_step inout GpOpt o inout GProg pg → 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.

opt_step takes the Opt inout, so its step counter ADVANCES — the frozen-Adam-t bug class (a scalar counter field on a by-value struct that never persists) 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

: Opt

: Opt {
    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_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 → i

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

@ opt_set_clip inout Opt o f maxn → v

@ opt_add inout Opt o GTape tp GVar p f alpha → v

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

@ _opt_gnorm Opt o GTape tp → f

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

@ opt_step inout Opt o GTape tp → v

One update step from the gradients currently on the tape.