checked.nu

gpukit/checked.nu — kernels a safe caller may write.

gk_run hands a kernel's source to the device compiler as is: a kernel can index any address it is given, so only unsafe code may launch one. A CHECKED kernel is written in a subset that cannot: its body reaches memory only through

LD(name, i) the element at i of buffer name (its last element when i is outside it) ST(name, i, v) store v there (nothing past its end; outputs only)

which gpukit defines per buffer from that buffer's own length, so every access is checked against the memory the call bound. The body sees its scalars by name, gid (this thread's index) and nthreads (how many run; a thread past it returns before the body). It may declare scalar locals, branch, loop and call the math functions below — nothing that makes or follows an address:

• no pointer declaration or dereference: a * that is not between two operands, a & that takes an address, [ ], ./-> (number literals aside), #, quotes and backslashes; • no name gpukit keeps for itself (anything starting with __, bar the rounding intrinsics), no asm, goto, typedef, struct, casts by keyword, new or delete; • calls only to LD / ST and a fixed list of math functions; • braces and parentheses that balance, so the body cannot close the kernel and write code after it.

gk_run_checked rejects a body outside the subset (F, and the reason in gk_checked_error) before anything is compiled. What remains is a kernel that cannot reach memory it was not given, so launching one is safe.

: ~ GkCall c ( gk_call_new ) ( gk_call_in_f c x xs ) // moved in; ( gk_call_take_f c y ) moves out ( gk_call_out_f c y ( vec_zeroed [f] n ) ) ( gk_call_i64 c n n ) : b ok ( gk_run_checked kit c scale2 ST(y, gid, 2.0 * LD(x, gid)); n 256 ) : ( Vec f ) ys ( gk_call_take_f c y )

API

: GkCall

: GkCall {
    ( Vec String ) names
    ( Vec i ) kinds
    ( Vec i ) slots
    ( Vec ( Vec f ) ) fv
    ( Vec ( Vec i ) ) iv
    ( Vec ( Slice f ) ) fs  // inputs bound as views: read where they are
    ( Vec ( Slice i ) ) is
    ( Vec GkBuf ) db  // device buffers (gk_call_dev / _dev_in): used where they live
    ( Vec GkHBuf ) hb  // 16-bit float buffers (gk_call_half_in): read where they live
    ( Vec GkBBuf ) qb  // byte buffers (gk_call_bytes_in): read where they live
    String err  // why gk_run_checked refused the kernel ("" when it did not)
    b exact  // gk_call_exact: the kernel is built for the buffers' lengths
}

One binding: kind 0 f64 input, 1 f64 output, 2 i64 input, 3 i64 output, 4 i64 scalar, 5 f64 scalar, 6 f64 input view, 7 i64 input view, 8 device buffer read, 9 device buffer read and written, 10 f64 output view, 11 i64 output view, 12 16-bit float device buffer read (as floats). slot indexes fv / iv for a buffer and holds the value (an f64's bits) for a scalar.

@ gk_call_new → GkCall

(the binding lists sized for a dozen at the start: a call is built on every launch, and growing them one push at a time was most of what a small kernel's launch cost on the host)

@ gk_call_in_f inout GkCall c s name sink ( Vec f ) v → v

@ gk_call_out_f inout GkCall c s name sink ( Vec f ) v → v

@ gk_call_in_i inout GkCall c s name sink ( Vec i ) v → v

@ gk_call_out_i inout GkCall c s name sink ( Vec i ) v → v

@ gk_call_view_f inout GkCall c s name ( Slice f ) v → v

An input read where it is — a view of the caller's data, no copy; the borrow checker keeps what it views alive while the call holds it.

@ gk_call_view_out_f inout GkCall c s name ( Slice f ) v → v

An output read and written where it is — a view of the caller's storage (a Slice is writable, as slice_put is): the kernel's stores land in it, and an element the kernel does not write keeps its value.

@ gk_call_view_out_i inout GkCall c s name ( Slice i ) v → v

@ gk_call_view_i inout GkCall c s name ( Slice i ) v → v

@ gk_call_i64 inout GkCall c s name i x → v

@ gk_call_f64 inout GkCall c s name f x → v

@ gk_call_dev_in inout GkCall c s name GkBuf b → v

A device buffer the kernel reads (gk_run_checked_dev): nothing is copied; the call holds another owner of the buffer's block while it holds the binding, and the caller keeps its own.

@ gk_call_half_in inout GkCall c s name GkHBuf h → v

A buffer of 16-bit floats (f16 or bf16) the kernel reads: LD(name, i) is the element as a float, converted exactly. Read only. (The call holds another owner of the buffer; the caller keeps its own.)

@ gk_call_bytes_in inout GkCall c s name GkBBuf q → v

A byte buffer the kernel reads (gk_run_checked_dev): LD gives a byte, LD16 / LD32 the aligned word holding it (see GkBBuf).

@ gk_call_bytes inout GkCall c s name GkBBuf q → v

A byte buffer the kernel reads and writes: as gk_call_bytes_in, and ST(name, i, v) stores the low byte of v at byte i (a store out of range is dropped).

@ gk_call_dev inout GkCall c s name GkBuf b → v

A device buffer the kernel reads and writes.

@ gk_call_exact inout GkCall c → v

Build the kernel for its device buffers' lengths (gk_run_checked_dev): each is a constant in the kernel's source and part of the key it is cached under — buffers of other lengths build another kernel. For a kernel whose dimensions are constants too (one compiled per shape): then a window, or a load whose index the loops bound, is seen to be in range at compile time, and its check folds away.

@ gk_call_set_i64 inout GkCall c s name i x → b

Give scalar name a new value — one call launched again and again. F when the call has no such scalar.

@ gk_call_set_f64 inout GkCall c s name f x → b

@ gk_call_take_f inout GkCall c s name → ( Vec f )

Move buffer name back out of the call (an empty Vec when there is none).

@ gk_call_take_i inout GkCall c s name → ( Vec i )

: GkcBuilt { String src String ename b shfl } // shfl: the blocks must be whole warps

Kernels checked and built before: key (name, bindings, body) → the built source and its entry point, shared (an Arc) so a launch reads them in place. A fit launches the same kernel every round; checking its body once is enough.

: GkcCache

: GkcCache {
    ( Vec String ) keys
    ( Vec ( Arc GkcBuilt ) ) built
    ( Vec String ) names  // each entry's kernel entry point
    ( HashMap String i ) ix  // key → entry
}

: ( Locked GkcCache ) g_gkc_cache ( __gkc_cache0 )

@ gk_checked_error GkCall c → s

Why the last gk_run_checked over c refused its kernel ("" when it did not).

: GkcU

: GkcU {
    ( Vec i ) id  // per token: its identifier's id, -1 for anything else
    ( Vec i ) var  // per id: 1 when it can name a variable
    ( Vec i ) per  // per id: 1 when per thread
    ( Vec i ) wr  // per id: 1 when a buffer the kernel writes
    ( Vec i ) kd  // per token: 1 if, 2 while, 3 for, 4 else, 5 {, 6 }, 7 a block-wide op, 8 LD, 0 other
    ( Vec i ) asg  // per token: 1 an assignment operator follows, 2 a ++ / -- beside it
    ( Vec i ) e1  // per token: if / while / for — its ) ; an assignment — its value's end
    ( Vec i ) e2  // per token: for — its first ;
    ( Vec i ) e3  // per token: for — its second ;
}

The uniformity walk works on ids: every identifier of the body interned once, with per-name flags — per thread, a written buffer, a name a value can be given to (not a type, not a callable) — so the question it asks of every token in every round is a lookup, not a search of a name list.

@ gk_launch_ok GkLaunch l → b

---

gpukit.nu

gpukit/gpukit.nu — the ergonomic GPU-compute facade over the gpu package.

The gpu package is the low-level interface: open a device, compile a CUDA-C kernel (JIT via NVRTC, or the host-C++ CPU backend), allocate device buffers, upload/download, build the argument vector, compute a grid, launch, sync, free. Every GPU-using package (onnx, objdet, anomaly, yoloe) hand- writes the same ~35 lines of that marshalling for each kernel it runs:

: GpuBuffer b_x ( gpu_alloc g ( n 8) ) // one per array … ( gpu_upload b_x (# u (vec_data [f] xs)) ) // upload each input … : (Vec i) args (vec_new [i]) // build the arg vector … ( vec_push [i] args (gpu_arg_buffer b_x) ) … ( gpu_launch k (gpu_grid n 256) 256 args ) ( gpu_sync g ) ( gpu_download (# *u (vec_data [f] out)) b_out )

gpukit collapses that to a list of typed bindings and one gk_run:

: GpuKit kit ( gk_open 0 ) : (Vec GkArg) call ( vec_new [GkArg] ) ( vec_push [GkArg] call ( gk_in_f xs ) ) // input double[] ( vec_push [GkArg] call ( gk_i64 n ) ) // long long scalar ( vec_push [GkArg] call ( gk_out_f out) ) // output double[] ( gk_run kit src my_kernel ( gk_grid n 256 ) 256 call )

gk_run compiles-and-caches the kernel by name, allocates a device buffer per buffer binding, uploads inputs, builds the arg vector in binding order, launches, syncs, downloads outputs, and lets every device buffer go. Kernel sources are cached on the kit, so a hot path compiles each kernel once.

gpukit adds no numerics of its own — it only marshals — so a kernel runs bit-for-bit the same through gk_run as through hand-written gpu_* calls, and the gpu package's CUDA / CPU-backend / pure equivalence is preserved.

Memory: a GpuKit is a handle on the kit's state (an Arc): every copy of it — a struct field, a Vec element, a capture — is the same kit, and its LAST owner releases the cached kernels, the device-memory pool and the device. Every GkBuf the kit hands out (dev.nu) holds the kit too, so the device outlives its memory whichever order the owners go in. Nothing is released by hand; gk_close is an optional early release of one owner.

API

: GkPoolG

: GkPoolG {
    i n  // rows, every kit's pools together
    b on
    i idle  // bytes currently held idle
    i max  // budget in bytes; <0 = not set yet, 0 = unlimited
    i clock  // LRU stamp source
}

The device-memory pool's process-wide counters and switches (see the pool section below). Every kit on every thread updates them, so they live behind one lock: as plain globals two threads' n + 1s could both read the same n, and the idle total that decides evictions drifted.

: ( Locked GkPoolG ) g_pool ( __gkp0 )

@ _gkp → GkPoolG

The totals now.

@ _gkp_add i dn i didle i tick → i

dn more rows and didle more idle bytes; the LRU clock advances by tick. The clock's new value.

: GkArg

: GkArg {
    i kind
    * u host  // host data pointer for a buffer binding
    i bytes  // buffer byte size
    i argval  // pre-encoded scalar arg (gpu_arg_*) for a scalar binding
}

A single argument to a kernel launch. kind 0 = input buffer 1 = output buffer 2 = scalar 3 = in/out buffer (uploaded, then copied back: what the kernel does not write keeps the host's values)

: GkKernelEntry

: GkKernelEntry {
    String name
    GpuKernel kernel
}

: GkInfo

: GkInfo {
    Gpu gpu
    b ok
    b census  // a census kit (gk_open_census): records, never launches
}

What a kit is and never changes once it is open: the device, whether it opened, whether it is a census kit. Shared, read without a lock.

: GkState

: GkState {
    ( Vec GkKernelEntry ) cache
    ( HashMap String i ) cix  // an entry name → its cache slot
    ( Vec i ) kcalls  // per cache slot: launches, counted only while gk_prof_on
    ( Vec i ) kns  // per cache slot: device time in those launches (CUDA events)
    ( Vec String ) csrc  // census only: each cache slot's kernel source
    ( Vec GpuBuffer ) pbuf
    ( Vec i ) pdptr
    ( Vec i ) pbytes
    ( Vec i ) pinuse
    ( Vec i ) pstamp
    GpuTimer ev0  // the profiler's event pair (gk_prof)
    GpuTimer ev1
    b async  // gk_async: launches on this kit do not wait for the device
}

What a kit changes as it runs: the kernel cache, the profiler's counts and event pair, and the device-memory pool. The pool is five parallel columns, one row per device block the kit allocated (in use by a GkBuf, or idle): the blocks themselves (owned — dropping a row frees its memory) and the plain columns a scan reads. Reached only through locked_take / locked_restore, so two threads using one kit take turns; nothing gpukit does while it holds the state hands control to code that could ask for it again.

% Drop GkState

% Drop GkState {
    @ drop GkState x → v {
        : ~ i idle 0
        : ~ i k 0
        ~ < k ( vec_len [i] . x pdptr ) {
            ? == ( __gk_col . x pinuse k ) 0 { = idle + idle ( __gk_col . x pbytes k ) } {}
            = k + k 1
        }
        : i _c ( _gkp_add - 0 ( vec_len [i] . x pdptr ) - 0 idle 0 )
    }
}

The process-wide pool counters (gk_pool_count / gk_pool_idle_bytes) sum every kit's pool; a kit that goes takes its rows out of them. Its fields — the pool's blocks, the cached kernels, the timers — are released by the compiler after this (drop glue).

: GpuKit

: GpuKit {
    ( Locked GkState ) st
    ( Arc GkInfo ) info
}

A GpuKit is a handle: every copy is the same kit, and the last owner releases it.

@ GpuKit_share GpuKit h → GpuKit

@ GpuKit_drop sink GpuKit h → v

@ _gk_ok GpuKit h → b

The kit's immutable part, lent.

@ _gk_census GpuKit h → b

@ gk_open i ordinal → GpuKit

Open device ordinal (CUDA when present, else the gpu package's CPU backend). The kit is a long-lived handle (its kernel cache and pool persist across calls). Check gk_ok; a failed open is still a kit.

@ gk_open_best → GpuKit

Open the device a caller who does not care should get: $NURL_GPU_DEVICE when set, else the highest compute capability with memory breaking ties. gk_open 0 binds ordinal 0, which on a box with an old card beside a new one is whichever the driver enumerates first — a 4 GB GTX 970 in front of a 24 GB RTX 4090, say, where a model that fits on one fails to allocate on the other. Prefer this everywhere the ordinal is not a user choice.

@ gk_none → GpuKit

A kit that opens nothing and is not ok (gk_ok F): what a slot holds once its kit has been let go.

@ gk_ok GpuKit kit → b

@ gk_gpu GpuKit kit → Gpu

The kit's device, for the gpu package's own calls (graphs, timers): another owner of it, so it stays valid however long it is kept.

@ gk_backend GpuKit kit → s

"cuda" or "cpu". A census kit answers "cpu": it takes the static backend's branches (see gk_open_census).

@ gk_jit GpuKit kit → b

T when this kit compiles kernels from source at run time — CUDA (NVRTC) or the CPU backend (host C++) — so a shape-specialised kernel, whose name and body are generated per call, costs one compile and nothing else. F on the static and WebGPU backends, whose kernel sets are fixed when the program is BUILT and looked up by entry name: a name carrying a run-time shape can never be in such a set, so a wrapper that specialises must ask this first and fall back to its generic kernel. F on a census kit, which records what a static build must link.

@ gk_open_census → GpuKit

── Kernel census ─────────────────────────────────────────────────────

A census kit opens no device and launches nothing. Every gk_run / gk_run_dev — so every gkd_* wrapper — validates its arguments exactly as it would on a device, then RECORDS the kernel it would compile (entry name and the exact source, once per name) and reports success. It answers gk_backend "cpu" and gk_jit F, which are the branches the gpu package's STATIC backend takes — so driving a program's kernel calls through a census kit yields precisely the kernel set a static build (a precompiled kernels_static.c, native or wasm32) has to link, from the same source builders that run on a device: nothing mirrored by hand, nothing to rot. Buffers for the calls come from gk_dbuf_new on the census kit itself: zeroed host memory the buffer owns, which no kernel reads.

@ gk_is_census GpuKit kit → b

@ gk_census_name GpuKit kit i k → String

What a census kit recorded: gk_kernel_count kernels, slot k's entry name and source ("" out of range).

@ gk_census_src GpuKit kit i k → String

@ gk_device_name GpuKit kit → s

Borrowed from the kit's device: valid while the kit is.

@ gk_bind_thread GpuKit kit → b

Make the kit's device current on the CALLING thread — see gpu_bind_thread. A single-threaded program never needs it; anything that hands device work to a pool or a fiber runtime does.

@ gk_mem_free GpuKit kit → i

Free / total device memory in bytes; 0 means the backend cannot say. CUDA asks the driver; the CPU backends report host RAM.

@ gk_mem_total GpuKit kit → i

@ gk_pool b on → v

@ gk_pool_enabled → b

@ gk_pool_count → i

Blocks the pools hold, in use or idle — for tests and for anyone wondering where the VRAM went.

@ gk_pool_idle_bytes → i

Bytes the pools are holding that nothing is using.

@ gk_pool_budget GpuKit kit i bytes → v

Cap the idle side at bytes (0 = unlimited). Trims immediately.

@ gk_pool_budget_bytes → i

@ _gk_pool_default GkInfo inf → v

The budget a caller who never sets one gets: $NURL_GK_POOL_MAX (bytes) when set, else a quarter of the device's memory — enough that a steady shape mix never evicts, small enough that a drifting one cannot eat the card. Resolved once, on the first allocation.

@ _gk_pool_take inout GkState kit i bytes → i

An idle block of exactly bytes, marked in-use — its device pointer, or 0.

@ _gk_pool_add inout GkState kit sink GpuBuffer gb → v

A fresh driver allocation joins the table, in use.

@ _gk_pool_give GpuKit h i dptr → v

The last owner of the block at dptr let go of it: idle in the pool, or — with the pool off — straight back to the driver.

@ gk_pool_release GpuKit kit → v

Hand every idle block back to the driver and drop it from the table. In-use blocks stay — this is a trim, not a reset.

@ _gk_pool_release inout GkState kit → v

@ gk_close sink GpuKit kit → v

Let go of kit now rather than at the end of its owner's scope: the device closes with the kit's last owner — which includes every GkBuf still alive from it.

@ gk_grid i n i block → i

Grid size for n threads at the given block size (re-export of gpu_grid).

@ gk_kernel_count GpuKit kit → i

Kernels in the kit's cache (each compiled once, by name).

@ gk_compile GpuKit kit s src s name → b

Warm the cache: compile src (entry name) into the kit now and report whether it succeeded, so a long-lived caller can detect a bad kernel at setup instead of on the first launch.

: s GK_EXPR_FNS x exp expf log logf sqrt sqrtf rsqrtf tanh tanhf erf erff fabs fabsf fmax fmaxf fmin fminf pow powf sin sinf cos cosf floor floorf ceil ceilf ``

@ _gk_ident_ch i c → b

@ gk_name_ok s name → b

A kernel or operation name: [A-Za-z0-9_]+.

@ gk_op_ok s op → b

@ gk_expr_ok s e → b

---

dev.nu

gpukit/dev.nu — device-RESIDENT buffers + dtype-aware kernels.

gk_run marshals host↔device around every call, which is right for one-shot compute but wrong for chained pipelines (an ndarray expression, a NN forward pass): the data should stay on the device between ops. This layer adds exactly that seam:

GkBuf an element-typed device allocation (GK_F32 | GK_F64 | GK_I64) gk_dbuf_new / free / upload / download (f64 host view, converting) gkdbuf_upload_i / downloadi (exact i64 host view, GK_I64 only) gk_run_dev cached-compile + launch + sync over RAW device args gkd_* ready-made dtype-generic kernels over GkBuf: elementwise (scalar + full N-D stride broadcast), unary maps, matmul, batched matmul, gather/scatter, row softmax, sum reduction

Numerics: GK_F64 kernels compute in double exactly like kernels.nu. GK_F32 buffers hold real float32 and kernels compute IN float32 (accumulation included) — true float32 semantics, matching numpy float32 / onnxruntime, NOT the host-tensor trick of f64-compute + grid-rounding. Uploads convert f64 → f32 via a staging buffer. GK_I64 holds long long (index tensors: gather/scatter/argmax); the f64 host view converts by C truncation, the _i entry points are exact.

API

@ _gk_p2 i a i b → i

Products of sizes as a kernel will form them: -1 when a factor is negative or the product overflows, which no buffer length equals.

@ _gk_p3 i a i b i c → i

@ _gk_p4 i a i b i c i d → i

@ gk_hbuf_ok GkHBuf h → b

@ gk_hbuf_len GkHBuf h → i

@ gk_hbuf_new GpuKit kit i n b bf → GkHBuf

A new buffer of n 16-bit floats (contents undefined until uploaded).

@ gk_hbuf_upload_bytes GpuKit kit GkHBuf h ( Slice u ) src → b

Upload n halves from a checked view of their bytes (2n of them, little endian, as a checkpoint stores them): F unless src holds that many.

@ gk_bbuf_ok GkBBuf q → b

@ gk_bbuf_len GkBBuf q → i

@ gk_bbuf_new GpuKit kit i n → GkBBuf

A new buffer of n bytes (contents undefined until uploaded).

@ gk_bbuf_upload_bytes GpuKit kit GkBBuf q ( Slice u ) src → b

Upload exactly the buffer's bytes from a checked view of them.

@ _gk_esz i dtype → i

@ _gk_tname i dtype → s

@ _gk_pfx i dtype → s

Kernel-name prefix per dtype so the cache never mixes element types.

@ _gk_scal i dtype f v → i

A scalar kernel argument holding value v in the buffer dtype's width: f32 bit pattern for GK_F32, f64 bits otherwise (the arg cell is 8 bytes; the kernel reads the leading 4 for a float parameter).

@ gk_buf_ok GkBuf b → b

@ gk_buf_len GkBuf b → i

@ gk_buf_dtype GkBuf b → i

@ gk_buf_esz GkBuf b → i

Bytes per element. Needed by anyone addressing a sub-range of a buffer (the device pointer is byte-addressed), which is why it is public rather than the file-private gkesz it wraps.

@ gk_dbuf_new GpuKit h i n i dtype → GkBuf

@ gk_dbuf_free sink GkBuf b → v

Let go of b now rather than at the end of its owner's scope.

@ gk_dbuf_upload_bytes GpuKit kit GkBuf b ( Slice u ) src → b

The same from a checked view of the bytes (a tensor's bytes in a mapped file, safetensor's st_tensor_bytes): F unless it holds the buffer's b.n elements. Safe, and as fast — one memcpy to the device.

@ gk_dbuf_download_bytes GpuKit kit GkBuf b ( Slice u ) dst → b

Device → a checked view of host bytes: F unless dst has room for the buffer's b.n elements. The download counterpart of gk_dbuf_upload_bytes.

@ gk_dbuf_download_host GpuKit kit GkBuf b GpuHost h → b

Device → page-locked host memory (a GpuHost from gpu_host_alloc): all of the buffer's bytes into the front of h. F unless h has room for them. A view of the buffer downloads part of it.

@ gk_dbuf_copy GpuKit kit GkBuf dst GkBuf src → b

Device → device: all of src into the front of dst. F unless both are there, share a dtype and dst holds at least src.n elements.

@ gk_arg_dev GkBuf b → i

Compile-cached launch over pre-built args (device pointers via gk_arg_dev, scalars via gpu_arg_i64/i32/f32). Syncs before returning unless autosync is off (below).

: ~ b g_gk_autosync T

Autosync: gk_run_dev (and every gkd_* kernel on top of it) normally syncs the device after each launch. An executor that chains hundreds of launches can turn this off, launch away, and gk_sync once at the end — the CUDA stream serialises kernels, and downloads synchronise implicitly; the CPU backend runs launches synchronously either way. gk_autosync switches it for every kit of the process; gk_async for one kit, which is what a library holding its own kit wants — it does not change how anyone else's launches behave.

@ gk_autosync b on → v

@ gk_async GpuKit h b on → v

@ gk_sync GpuKit kit → b

: ~ b g_gk_prof F

── Per-kernel profiling ────────────────────────────────────────────── "The model is slow" is not actionable; "78% of the frame is in three kernels" is. With profiling on, every gk_run_dev launch is bracketed by a CUDA event pair (the kit's own, made on first use) and the device time is accumulated into the kernel's cache slot, so gk_prof_report ranks the kernels by where the time actually went. Events measure the GPU, not the launch — a host clock around an async launch measures neither.

It is not free: each launch waits for its own end event, which serialises the stream. Turn it on to find the hot kernel, off to time the program.

@ gk_prof GpuKit h b on → v

@ gk_prof_on → b

@ gk_prof_reset GpuKit h → v

Drop every accumulated count — e.g. after a warm-up frame, so the numbers describe the steady state and not the compiles.

@ gk_prof_total GpuKit h → i

Total device time accumulated so far, in nanoseconds.

@ gk_prof_report GpuKit h → v

Kernels by device time, slowest first, with the share of the total.

@ gkd_ew GpuKit kit s opname s op GkBuf o GkBuf a GkBuf b → b

op is one of + - * / and opname an identifier (gk_op_ok, gk_name_ok); F otherwise.

@ gkd_add GpuKit kit GkBuf o GkBuf a GkBuf b → b

@ gkd_sub GpuKit kit GkBuf o GkBuf a GkBuf b → b

@ gkd_mul GpuKit kit GkBuf o GkBuf a GkBuf b → b

@ gkd_div GpuKit kit GkBuf o GkBuf a GkBuf b → b

@ _gk_vi ( Vec i ) v i k → i

@ gkd_ew_bc GpuKit kit s opname s op GkBuf o GkBuf a GkBuf b ( Vec i ) odims ( Vec i ) astr ( Vec i ) bstr → b

out over odims; a and b read through their stride tables (entries ≥ 0, 0 = broadcast). Float dtypes only.

@ _gkd_dim_arg inout GkCall c s p i k i v → v

Bind scalar <p><k> (d0, A3, …) to v.

@ gkd_map GpuKit kit s mapname s expr GkBuf o GkBuf a → b

expr is checked like gk_map_f's (gk_expr_ok); F otherwise.

@ gkd_relu GpuKit kit GkBuf o GkBuf a → b

@ gkd_sigmoid GpuKit kit GkBuf o GkBuf a → b

@ gkd_exp GpuKit kit GkBuf o GkBuf a → b

@ gkd_tanh GpuKit kit GkBuf o GkBuf a → b

@ gkd_sqrt GpuKit kit GkBuf o GkBuf a → b

@ gkd_log GpuKit kit GkBuf o GkBuf a → b

: i GKD_RT 8

Rows / columns of C one CPU-backend thread owns. See gktile_body.

: i GKD_CT 32

@ _gkd_ceil i x i q → i

@ _gkd_mac i dtype s dst s x s y → String

The multiply-accumulate step, spelled the way each element type's contract requires: F64 with explicit round-to-nearest intrinsics (NVRTC's fmad contraction would otherwise fuse and change the bits), F32 plain (its contract is true-float32 semantics, which the verified model goldens pin). dst and the product operands are C expressions.

: i GKD_BM 64

The shared-memory tiled matmul for the CUDA backend — the body behind gkd_gemm and gkd_bmm both (gkdsmem_cbody).

One thread per output element (the shape this file shipped with) is bandwidth-bound: every MAC costs a fresh global load of B, so a 4090 runs it at ~3 TFLOP/s, about 4% of what its FP32 units can do. Here a 256-thread block owns a 64x64 tile of Y, stages 64x16 of A and 16x64 of B into shared memory, and each thread keeps a 4x4 register tile. B is read once per tile instead of once per output row, and the arithmetic intensity goes from 1 MAC per load to 32. Measured on a 4090 at 783x1024x1024 f32: 3.1 -> 11.5 TFLOP/s, and 0.6 -> 11.1 with B transposed (the per-element kernel reads a transposed B along the wrong axis entirely).

The accumulation order is unchanged: each output still sums t = 0..K-1 ascending, only blocked. Out-of-range staging slots hold a zero, so a ragged tail adds exact 0.0 terms — the value is bit-identical to the per-element kernel, which is what the verified model goldens pin. Nothing here reassociates.

transb is baked into the body (and into the kernel name) rather than branched on per k: with B as [N,K] the staging loop walks k contiguously per column, so transb=1 loads coalesced too — which is why, unlike the CPU path, this one does not care which layout the weight arrives in.

batched picks the bmm signature (per-batch A/B strides as/bs, a bare store) over the gemm one (alpha/beta and a per-column bias). The tile: a block owns GKD_BM x GKD_BN of Y and each of its 256 threads a GKD_TM x GKD_TN register block, stepping K in GKD_BK.

Why 128x64x16 with an 8x4 register tile and not the obvious 64x64 with 4x4: at 4x4 a thread reads 8 floats from shared per k step to do 16 MACs, and Ada's 128 B/clk of shared bandwidth cannot keep 128 FMA lanes fed at that ratio — the kernel stalls on shared, not on math. 8x4 halves the ratio to 12 reads per 32 MACs. Going the whole way to 128x128/8x8 halves it again, but M here is ~800 rows and N is often 1024, which is 56 blocks — under half of a 4090's 128 SMs, so the wider tile loses more to idle SMs than it gains per SM.

: i GKD_BN 64

: i GKD_BK 16

: i GKD_TM 4

: i GKD_TN 4

@ _gkd_in inout GkCall c s name GkBuf b GkBuf out → v

Bind b as an input of c: read-only (restrict, and the read-only cache) unless it overlaps out, the buffer the kernel writes — an op called in place reads through a writable binding, as before.

@ _gkd_ab i dtype s nm → String

alpha / beta as the element type: a checked scalar is a double, and an f32 GEMM's epilogue rounds in float.

@ _gkd_smem_cbody s tn i dtype b transb b batched b big → String

The CUDA shared-memory tile, checked (see the notes above for the design and the measurements: a block owns a 64x64 or 128x64 tile of Y, stages 16-wide slices of A and B into shared memory, each thread keeps a 4x4 or 8x4 register block). The tiles are SHARED arrays — an f32 one read four at a time with LD4, a 16-byte shared load — and every index into them and into the register block is made of tid, loop counters and constants, so its check folds. The staging loads are checked against their buffers. Each output still sums t = 0..K-1 ascending, out-of-range slots adding exact zeros: bit-identical to the raw tile.

Bindings A, B, Y (and Cb, the bias, for the GEMM); scalars M, N, K, and alpha, beta, hasB (GEMM) or as, bs (batched: per-batch strides, 0 broadcasting).

@ gkd_scatter GpuKit kit GkBuf y GkBuf d GkBuf ix GkBuf u i outer i ax i inner i nidx → b

y = d with y[outer, ix[g], inner] = u[outer, g, inner]. When several indices collide the surviving write is unspecified (ONNX leaves this undefined too). y may be d itself for in-place update.

@ gkd_softmax_rows GpuKit kit GkBuf o GkBuf a i rows i cols → b

@ gkd_sum GpuKit kit GkBuf a → ?f

---

kernels.nu

gpukit/kernels.nu — a small library of ready-made f64 kernels over gk_run.

These cover the operations ML code writes over and over. They generate the CUDA-C source, cache it on the kit (by a fixed entry name), and marshal through gk_run — so a caller never writes a kernel or a device buffer for the common cases.

Numerics: f is a C double, so every buffer is float64. Elementwise ops and gk_matmul_f accumulate in the same order a sequential host loop would, so they are bit-identical across the CUDA / CPU / pure backends and to a naive host implementation. gk_reduce_sum_f / gk_dot_f combine per-thread partial sums (a parallel order), so they match a host sum to rounding but are not guaranteed bit-identical to a strictly sequential accumulation.

API

@ gk_add_f GpuKit kit ( Vec f ) out ( Vec f ) a ( Vec f ) b → b

@ gk_sub_f GpuKit kit ( Vec f ) out ( Vec f ) a ( Vec f ) b → b

@ gk_mul_f GpuKit kit ( Vec f ) out ( Vec f ) a ( Vec f ) b → b

@ gk_div_f GpuKit kit ( Vec f ) out ( Vec f ) a ( Vec f ) b → b

@ gk_map_f GpuKit kit s kname ( Vec f ) out ( Vec f ) in s expr → b

── Unary map: out[i] = <expr>, with x bound to in[i] ──────────────── kname is the CUDA entry name (a valid C identifier, stable per expr so caching works). expr is C over the local double x, e.g. x*x, 1.0/(1.0+exp(-x)), x>0.0?x:0.0 — checked (gk_expr_ok): F for anything else, and for a name that is not an identifier.

: i GK_MM_RT 8 // rows of C per thread (CPU tile)

: i GK_MM_CT 32 // cols of C per thread (CPU tile)

@ _gk_cmac i dtype s idx s x s y → String

A multiply-accumulate into per-thread array acc at idx, spelled as each element type's contract requires: f64 with explicit round-to-nearest (no fused multiply-add), f32 plain (true float32 semantics; CUDA may fuse it, as it did the raw kernels').

@ _gk_tile_body s tn i dtype b batched b gemm → String

The CPU backend's register-tiled matmul body, checked: a thread owns an 8x32 block of Y and reads B along a row — 32 contiguous elements reused by all 8 rows — so the host compiler vectorises the inner loop. A's rows and B's column strip are 2-D windows, opened once per thread; the whole-tile loop runs to constant bounds, so no load in it carries a check. Per output the sum still runs t = 0..K ascending: every value is bit-identical to a per-element kernel and to a sequential host loop.

Bindings A (M×K), B (K×N) and Y (M×N), scalars M, K, N. batched: Y is batches of M×N, A and B advance by scalars as / bs per batch (0 broadcasts), nb batches. gemm: the epilogue is alpha·acc + beta·bias with a per-column bias Cb read when hasB.

@ gk_dims_fit_len i len i r i c → b

Whether a buffer of len elements holds an r×c matrix: neither dimension negative, and r·c ≤ len without the product overflowing.

@ _gk_zeros i n → ( Vec f )

@ _gk_partial_threads i n → i

@ gk_reduce_sum_f GpuKit kit ( Vec f ) x → ?f

Sum of x. None on a device error.

@ gk_dot_f GpuKit kit ( Vec f ) a ( Vec f ) b → ?f

Dot product a·b over a's length. None on a device error, or when b is shorter than a.

---

devbuf.nu

gpukit/devbuf.nu — the device buffer type (GkBuf) and the owner of its block, on their own so that dev.nu and checked.nu can both name them before either's functions (dev.nu's ops run as checked kernels, and a checked call binds GkBufs).

API

: i GK_F64 0

: i GK_F32 1

: i GK_I64 2

% Drop GkMemImpl { @ drop GkMemImpl x → v { ? == 0 ( vec_len [u] . x host ) { ( _gk_pool_give . x kit . x dptr ) } {} } }

: GkMem { ( Arc GkMemImpl ) a }

@ GkMem_share GkMem h → GkMem

@ GkMem_drop sink GkMem h → v

---

devops.nu

gpukit/devops.nu — the NN operator family over device-resident GkBufs.

These are the kernels a CNN / transformer forward pass needs beyond the arithmetic core in dev.nu: gemm, conv/pool, normalisation, activation, axis softmax, axis data movement (concat/slice/permute/resize/expand), reductions and argmax/index selection. The kernel bodies are lifted from packages/onnx's original f32 operator set and generalised over the element type, keeping the arithmetic ORDER identical — so the GK_F32 instantiation is bit-for-bit the kernels the onnx/objdet/yoloe models were verified with, and GK_F64 gets the same operators in double.

Conventions:

• every wrapper validates buffers, dtypes and sizes and fails closed

(returns F) — a wrong shape never reaches the device;

• layout metadata is explicit (NCHW dims, axis views (outer,ax,inner));

batch N is assumed 1 for the conv family, group 1;

• arithmetic ops require a float dtype; pure data movement

(copy/slice/permute/resize/expand) accepts GK_I64 too;

• index outputs (argmax) are GK_I64.

API

@ gkd_gemm GpuKit kit GkBuf y GkBuf a GkBuf b GkBuf c i hasb i m i n i k f alpha f beta i transb → b

@ gkd_conv2d GpuKit kit GkBuf y GkBuf x GkBuf w GkBuf bias i hasb i cin i h i wd i cout i kh i kw i oh i ow i ph i pw i sh i sw → b

@ gkd_conv2d_dil GpuKit kit GkBuf y GkBuf x GkBuf w GkBuf bias i hasb i cin i h i wd i cout i kh i kw i oh i ow i ph i pw i sh i sw i dh i dw → b

── 2-D DILATED convolution, NCHW, batch 1, group 1 ────────────────── The atrous variant (U²-Net's RSU4F, DeepLab-family heads): tap (r, s) reads input row iy = oy·sh − ph + r·dh. gkd_conv2d is the dh=dw=1 special case and keeps its specialised fast kernel; this one is a plain one-thread-per-output body, because dilated layers in the models that need them are a handful of small feature maps, not the hot path.

@ gkd_convtranspose2d GpuKit kit GkBuf y GkBuf x GkBuf w GkBuf bias i hasb i cin i h i wd i cout i kh i kw i oh i ow i ph i pw i sh i sw → b

@ gkd_maxpool2d GpuKit kit GkBuf y GkBuf x i c i h i wd i kh i kw i oh i ow i sh i sw i ph i pw → b

@ gkd_batchnorm GpuKit kit GkBuf y GkBuf x GkBuf sc GkBuf bb GkBuf mean GkBuf var i c i hw f eps → b

@ gkd_leakyrelu GpuKit kit GkBuf y GkBuf x f alpha → b

@ gkd_clip GpuKit kit GkBuf y GkBuf x f lo f hi → b

@ gkd_erf GpuKit kit GkBuf y GkBuf x → b

@ gkd_layernorm GpuKit kit GkBuf y GkBuf x GkBuf sc GkBuf bi i outer i ax f eps → b

: i GKD_ATTN_SMEM 33792

Shared-memory budget one attention block may take: a third of the 100 KB an Ada SM can carve out, so three blocks stay resident.

@ gkd_attention_ok GpuKit kit i hd → b

Whether gkd_attention will run for this head width on this backend — so a caller can decide, before it allocates, whether it needs the [heads, n, nkv] score buffer the composed fallback would want. At a 72-frame window of 783 tokens that buffer is 2.8 GB, and it is the difference between a long sequence fitting and not.

@ gkd_attention GpuKit kit GkBuf o GkBuf q GkBuf k GkBuf v

Unmasked: every key is visible to every query.

@ gkd_attention_masked GpuKit kit GkBuf o GkBuf q GkBuf k GkBuf v GkBuf mask

With a per-key additive bias: mask holds (heads/hpb) rows of nkv floats, head h reading row h/hpb, and the bias is added to the score before the softmax — 0 keeps a key, -1e30 drops it. That is how a padded batch says "these positions are not there": the key never reaches the running maximum and contributes exactly zero to the denominator, so the row is the row it would have been without the padding. Pass a zero GkBuf (dptr 0) for mask to run unmasked — the kernel is then character-for-character the one gkd_attention has always generated, so nothing already compiled is invalidated.

@ gkd_attention_batch GpuKit kit GkBuf o GkBuf q GkBuf k GkBuf v GkBuf mask

The fused attention over a BATCH of sequences in the layout a transformer's linear layers actually produce: Q, K, V and O are [batch·n, heads·hd] — token rows, heads interleaved within the row — NOT the [heads, n, hd] gkd_attention wants. A caller running batch=1 used two gkd_perm round trips per tensor per block to get in and out of head-major layout; a caller running a real batch would also need a 4-D permute whose dims gkd_perm bakes into the kernel name, i.e. an NVRTC compile per (batch, length) pair. This entry removes both: the kernel reads the interleaved layout directly (row stride heads·hd, head offset h·hd — the loads along d stay contiguous and coalesced exactly as before), and no permuted copy of Q/K/V/O ever exists.

mask is one additive row of nkv floats PER BATCH ELEMENT ([batch, nkv]; 0 keeps a key, -1e30 drops it), or a zero GkBuf to run unmasked. Same tile, same online softmax, same accumulation order as gkd_attention_masked — a sequence embedded at batch b is bit-identical to the same sequence permuted through the head-major entry. gkd_attention_ok answers for this entry too (the tile constraint is the same).

@ gkd_expandlast GpuKit kit GkBuf y GkBuf x i outer i rep → b

Expand the last axis: input (outer,1) → (outer,rep), Y[o,r] = X[o].

@ gkd_reducel2 GpuKit kit GkBuf y GkBuf x i outer i ax → b

L2 norm along the last axis: input (outer, ax) → output (outer).

@ gkd_argmax GpuKit kit GkBuf y GkBuf x i outer i ax → b

ArgMax along the last axis: input (outer, ax) in any element type → GK_I64 indices (outer). Ties resolve to the first maximum.

@ gkd_eos_gather GpuKit kit GkBuf y GkBuf data GkBuf tok i bsz i l i d → b

EOS read-out (CLIP): given per-row token ids tok [B,L] (GK_I64) and features data [B,L,D], select each row's max-id token's features → Y [B,D]. One thread per output element.

@ gkd_rope2d GpuKit kit GkBuf x GkBuf rows GkBuf cols GkBuf cosb GkBuf sinb

@ gkd_rope3d GpuKit kit GkBuf x GkBuf fr GkBuf rw GkBuf cl GkBuf cosb GkBuf sinb