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_drop sink GTape h → v

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

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

---

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_drop sink GpPlan h → v

: ~ i g_gpm_next 1

@ _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_in String o * GProgImpl pg i inid i C → v

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

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

@ _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_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_fill s kname String o → v

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

@ _gpf_smax → i

@ _gpf_sel String o i id s ix → v

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

@ _gpf_sgel String o i id s ix → v

: 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_drop sink GpFuse h → v

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

---

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

: 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_drop sink GProg h → v

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

: 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_drop sink GpOpt h → v

@ _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_ew_block ( Vec i ) meta ( Vec i ) ost ( Vec i ) xe ( Vec i ) ye → v

An ew block [nd ost xe ye].

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

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

---

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_drop sink Opt h → v

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