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 ) ( gpu_free b_x ) … // free each buffer

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 ) ( vec_free [GkArg] 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 frees every device buffer. 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: GpuKit carries a stable kernel-cache Vec, so it is passed by value and mutated across calls (like an ArgParser). Free it with gk_close, which releases the cached kernels and the device.

API

: 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

: GkKernelEntry

: GkKernelEntry {
    String name
    GpuKernel kernel
}

: GpuKit

: GpuKit {
    Gpu gpu
    b ok
    ( Vec GkKernelEntry ) cache
}

@ gk_open i ordinal → *GpuKit

Open device ordinal (CUDA when present, else the gpu package's CPU backend). Returns a heap *GpuKit so it can be held as a long-lived singleton (its kernel cache persists across calls). Check gk_ok; even a failed open is safe to gk_close.

@ gk_ok * GpuKit kit → b

@ gk_backend * GpuKit kit → s

"cuda" or "cpu".

@ gk_device_name * GpuKit kit → s

@ gk_close * GpuKit kit → v

Release every cached kernel, close the device, and free the kit.

@ gk_grid i n i block → i

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

@ gk_in_f ( Vec f ) v → GkArg

@ gk_in_i ( Vec i ) v → GkArg

@ gk_out_f ( Vec f ) v → GkArg

Output buffers must be pre-sized by the caller; results are copied back in.

@ gk_out_i ( Vec i ) v → GkArg

@ gk_buf_in * u host i bytes → GkArg

Raw buffers, for element layouts other than f64/i64 (e.g. a packed float32 host buffer): pass the host pointer and byte size directly.

@ gk_buf_out * u host i bytes → GkArg

@ gk_i64 i v → GkArg

@ gk_i32 i v → GkArg

@ gk_f32 f v → GkArg

@ _gk_get_kernel * GpuKit kit s src s name → GpuKernel

Compile src (entry name) once per kit; subsequent calls with the same name reuse the cached kernel. A failed compile returns a not-ok kernel and is not cached (so a fixed source can be retried).

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

@ gk_run * GpuKit kit s src s name i grid i block ( Vec GkArg ) call → b

Compile-cached, marshal, launch, sync, download, free. call lists the kernel's arguments in declaration order (buffers and scalars interleaved exactly as the kernel signature expects). Returns F on any device error.

---

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

: i GK_F64 0

: i GK_F32 1

: i GK_I64 2

: GkBuf

: GkBuf {
    i dptr
    i n
    i dtype
}

An element-typed device allocation. dptr 0 = failed (safe to free).

@ _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_dbuf_new * GpuKit kit i n i dtype → GkBuf

@ gk_dbuf_free GkBuf b → v

@ gk_dbuf_upload * GpuKit kit GkBuf b ( Vec f ) src → b

Host f64 vector → device (converted to the buffer's element type: f32 rounds, i64 truncates like a C cast). Copies min(len(src), b.n) elements; missing elements upload as 0. The direct (non-converting) f64 path is taken only when src actually holds b.n elements — a short vector goes through the padding stage, never out of bounds.

@ gk_dbuf_download * GpuKit kit GkBuf b ( Vec f ) dst → b

Device → host f64 vector (converted from the buffer's element type). dst must already hold at least b.n elements (fails closed otherwise); the first b.n are overwritten in place.

@ gk_dbuf_upload_i * GpuKit kit GkBuf b ( Vec i ) src → b

Exact i64 host view (GK_I64 buffers only — index tensors must never make a round trip through f64). Same length contracts as the f64 view.

@ gk_dbuf_download_i * GpuKit kit GkBuf b ( Vec i ) dst → b

@ 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 b on → v

@ gk_sync * GpuKit kit → b

@ gk_run_dev * GpuKit kit s src s name i grid i block ( Vec i ) args → b

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

@ 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_map * GpuKit kit s mapname s expr GkBuf o GkBuf a → b

@ 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

@ gkd_matmul * GpuKit kit GkBuf c GkBuf a GkBuf b i m i k i n → b

@ gkd_bmm * GpuKit kit GkBuf y GkBuf a GkBuf b i batch i m i kk i n i abatch i bbatch → b

@ gkd_gather * GpuKit kit GkBuf y GkBuf d GkBuf ix i outer i axin i inner i nidx → b

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

@ gk_matmul_f * GpuKit kit ( Vec f ) c ( Vec f ) a ( Vec f ) b i m i k i n → b

── Matrix multiply: C[M×N] = A[M×K] · B[K×N] (row-major) ────────────── Each output element sums t=0..K in order, exactly as a host loop — so this is bit-identical to a naive sequential matmul.

@ _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 (equal lengths assumed). None on a device error.

---

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

@ gkd_softmax_ax * GpuKit kit GkBuf y GkBuf x i outer i ax i inner → b

@ gkd_copy_ax * GpuKit kit GkBuf dst GkBuf src i outer i src_ax i inner i dst_ax i off → b

Copy src into a concat output along an axis viewed as (outer, src_ax, inner); the slot starts at off in the output's axis of size dst_ax.

@ gkd_slice_ax * GpuKit kit GkBuf dst GkBuf src i outer i sz i inner i src_ax i soff → b

Extract a contiguous axis slice into a fresh tensor: viewing src as (outer, src_ax, inner), dst[o,a,i] = src[o, soff+a, i] for a in [0,sz).

@ gkd_perm * GpuKit kit GkBuf y GkBuf x ( Vec i ) dims ( Vec i ) perm → b

General N-D (≤6) transpose: output axis k reads input axis perm[k]. dims are the INPUT dims; both vecs hold ndim (≤6) entries and perm must be a permutation of 0..ndim−1.

@ gkd_resize_nn * GpuKit kit GkBuf y GkBuf x i c i h i wd i oh i ow i sh i sw → b

Nearest-neighbour upsample by integer factors (NCHW): Y[c,oy,ox] = X[c, oy/sh, ox/sw].

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