Kernels from Rust
Imported GPU kernels are Rust functions, type-checked and borrow-checked by rustc. The same source produces SPIR-V for wgpu and raw Vulkan, PTX for NVIDIA, and an AMDGCN code object for ROCm, not separate implementations for each backend. (Other kernels are synthesized; see Fusion.)
A kernel looks like this
#[kernel]
pub fn elementwise_add(a: &[f32], b: &[f32], out: &mut [f32]) {
let i = thread_index();
if i < out.len() {
out[i] = a[i] + b[i];
}
}
The #[kernel] attribute marks the function for import by pootc. The ABI is slices:
inputs as &[T], outputs as &mut [T]. The attribute renames the function into the reserved
__poot_kernel_<name> namespace for pootc; a normal host build gets a typed launch wrapper
instead of the body, so kernel bodies do not need to be host-compatible.
The import pipeline
Step through each stage to see how a #[kernel] fn becomes a GPU binary:
pootc is a thin rustc driver. It calls rustc normally (producing the normal compiled
artifact) and hooks after_analysis to visit every #[kernel] fn's MIR body using
Stable MIR 1, the stable, versioned interface to rustc's Mid-level IR
that does not change with every nightly.
Kernel IR (poot-kernel-ir) is a small CFG-based IR: basic blocks with statements
(Assign, WorkgroupLocalWrite, VectorStore, and the Wmma* matrix ops) and terminators
(Goto, SwitchInt, ThreadIndexCall, Barrier, Return, Trap). It
decouples the importer from the codegen backends, the importer only needs to know the
kernel IR; the backends only need to consume it. Rustc's bounds check on a slice index (a[i])
imports as a Trap; what a trap does on each target is under Failed checks.
How kernels connect to the graph executor
A kernel binary does not execute on its own; it is bound to a compute Plan produced by
plan_eqn, which connects the import pipeline to the graph planning layer. compile runs
this planning as its last step. A plan body is either an imported kernel (the committed
assets below) or a synthesized one; the planner chooses by region description and target, never by
model, and Fusion and Planning owns which ops take which.
Imported kernel sources live under crates/pootc/kernels/<family>/, one subdirectory per family:
contraction, packed and sampling, plus a probe family for kernels a backend crate loads directly.
Accepted Rust subset
This section is the one statement of what a kernel may contain. It is derived from the importer
(crates/pootc/src/import.rs) and the poot_kernel_intrinsics crate. Anything outside it produces a
named unsupported in kernel: ... diagnostic from pootc, never a silent miscompile.
Types. bool, usize, i32, u32, f32, f64, f16, the unit type, references and raw pointers
to those, and slices of them. Structs, tuples, arrays, enums, closures, trait objects and heap types
(Vec, Box, String) have no lowering and are rejected.
Expressions. Integer and float arithmetic (+, -, *, /, %), bit operations and shifts,
comparisons, unary - and !, numeric as casts, slice indexing and len(), and if / else.
Loops. while; for k in a..b and for k in a..=b; and for x in slice, which reads
slice[counter] for *x. .rev() and .step_by() are rejected.
Float and integer methods. On floats, sqrt, abs, floor, ceil, trunc, round, exp, ln,
sin, cos; max and min on floats and on integers; and f32::from_bits(u32), which reinterprets a
u32 as an f32 (for example a scale packed into a metadata buffer). The importer matches the
core/std method by its resolved path, so a kernel-authored helper with the same name is an ordinary call
and is rejected.
Intrinsics. A call into the poot_kernel_intrinsics crate is the only call a kernel may make besides the
methods above. The importer recognizes an intrinsic by the callee's resolved declaration, not its bare name,
so a kernel never redeclares one.
| Intrinsic | Meaning |
|---|---|
thread_index() | The global dispatch index along X |
local_index() | The lane within its workgroup along X |
group_index(), group_index_y() | The workgroup index along X, Y |
workgroup_barrier() | Workgroup barrier; earlier workgroup-local writes are visible to every lane after it |
wg_write(array, i, v), wg_read(array, i) | Workgroup-local (LDS) f32 access; array must be a literal array id |
atomic_add(buffer, i, v) | Atomic u32 add on buffer[i], returning the old value; buffer is a slice parameter and i a usize local, not a literal |
no_contract_add(a, b) | a + b that no backend may contract with a producing multiply into an FMA |
const WORKGROUP_SIZE: usize = N; in the kernel sets the workgroup lane count (default 64).
Not supported. A call to any other function is rejected, so helper functions are not inlined and
unwrap, expect, recursion and standard-library I/O have no lowering. Private local arrays with dynamic
indices are rejected as well; use workgroup memory (below).
Failed checks
A MIR Assert (a slice bounds check, an arithmetic overflow, a divide or shift check) lowers to a branch that
traps on failure; Unreachable lowers directly to a trap. A bounds check that a dominating guard provably
repeats is dropped, and a check no guard covers stays a trap.
| Target | A trap |
|---|---|
| ROCm, PTX | A device trap instruction |
| wgpu, raw Vulkan | SPIR-V has no trap. poot-codegen accumulates a per-thread bit for the failed site (up to 32 sites per kernel; further sites share the last bit) into the dispatch's error word and the thread carries on under robust buffer access. The executor discards the dispatch and raises a typed fault at its next sync point |
| XDNA2 NPU (AIE) | poot-codegen refuses a body that contains a trap with a typed error |
LDS arrays and barriers
Workgroup-local (shared) memory is written and read with wg_write and wg_read, indexed by
a literal array id:
pub fn __poot_kernel_wg_reduce(a: &[f32], out: &mut [f32]) {
let lane = local_index();
wg_write(0, lane, a[lane]); // array 0, slot `lane`
workgroup_barrier();
if lane == 0 {
let mut acc = 0.0f32;
for j in 0..64 { acc = acc + wg_read(0, j); }
out[0] = acc;
}
}
Use LDS instead of a local [f32; N] for SPIR-V: private arrays with dynamic indices
crash some driver stacks (RADV on AMD). The generated attention kernel
keeps its running output in an LDS array for this reason. LDS works on wgpu, raw Vulkan, ROCm and PTX.
The structurization pass converts the kernel IR's unstructured CFG to structured control flow (proper loops, if/else) before SPIR-V emit.
.kir.json assets and drift guarding
Each imported kernel has a committed .kir.json under its consuming crate's assets/
(poot-graph-plan/assets/ for every kernel the planner chooses between; poot-rocm-device/assets/ and
poot-wgpu-device/src/tests/assets/ for the handful a backend crate loads directly), generated from a source
under crates/pootc/kernels/<family>/. One manifest table lists every asset, its source and its
destination; a test regenerates each asset from its source and checks it byte-for-byte against the
committed one, and separately checks that each destination directory holds exactly its manifest's
entries. just regen-kernel-assets regenerates all of them.
Committing the kernel IR keeps GPU binaries reproducible from the repo without running pootc on every build.
Verification
Before codegen, Body::verify checks every matrix-fragment operation against the declared
Ty::Fragment type of its locals: the role (A, B or accumulator), the operand dtype and the tile
shape must match at every producer and consumer, and a fragment cannot be assigned like a scalar.
A mismatch returns a typed error with the offending statement location. Matrix operations are
collective over one 32-lane subgroup, so the verifier also rejects any matrix operation reachable
under a branch that can split a subgroup (DivergentFragmentParticipation); selecting a whole
subgroup, such as warp 0, is accepted. The accumulator is an ordinary variable on every target, so
it carries across a K loop; target and shape restrictions still apply.
committed_assets_match_their_kernel_sources regenerates every committed asset from its source and
fails if any is stale. Executor tests then dispatch the kernels on each device backend and compare against
the CPU oracle under the contract described in Execution. PTX tests run on
rented NVIDIA pods.