Introduction

This book documents the internals of zyx — a machine learning library and compiler.

Who This is For

This book is for developers who want to understand how zyx works under the hood, namely the architecture decisions.

Two Execution Modes

Zyx operates in two modes:

• Eager (default): Tensor operations fuse into kernels as you write them. When fusion is no longer possible, the pending kernel is compiled and executed when its value is needed. No separate realize step. Ideal for one-off work like data preprocessing and initialization.

• Tape: Wrap loops in a Tape to enable lazy graph building, autograd, and complex optimization (fusion variants compared by measured time, device placement, plan caching across structurally identical iterations). Think of it as torch.compile, but less strict. Kernels are shared across both modes.

The Architecture at a Glance

Every tensor operation is appended somewhere; what happens next depends on the mode:

Eager — direct fusion, no graph:

Tensor op ──► append to kernel ──► compile + execute

Each op is appended directly to the kernel that produced its inputs. When a new op isn’t fused, the pending kernel compiles and executes. No graph, no separate realize step. Everything is automated and runs in the background, user can’t really observe it and shouldn’t care about it.

In eager mode, the fusion heuristic is conservative, won’t overfuse as to not break the premise of the eager mode.

Tape — lazy graph, egraph exploration:

Tape is used for repeated computation. The opset and kernels are the same, it’s just when you use tape, you tell zyx that you are fine with being explicit (manual realize) and deferred (computation-wise) while you want to reap benefits of more optimizations (which comes from better fusion due to deferred computtation) as well as autograd. Additionally zyx takes this as a signal that it can take more time optimizing the DAG (directed acyclic graph) of tensor ops.

Tensor op ──► graph op (accumulates) ──► Autograd (reverse-mode on same graph)
                                  │
                                  ▼
                           Tape::realize()
                                  │
                                  ▼
                    Kernelizer (batch fusion variants
                     + device placement)
                                  │
                                  ▼
                    Kernel IR ──► Opt ──► Codegen ──► Execute

Inside a tape, graph ops accumulate lazily. At realize time, the kernelizer builds fusion variants for the full graph, each candidate kernel is autotuned and measured, and extraction keeps the fastest path (including device placement). Autograd reuses the same graph ops — the tape prevents their deletion so the backward pass can traverse them.

The trade-off in tape mode is that evaluation is lazy — you must call realize() to trigger computation. But this laziness enables optimizations that eager execution cannot do: autograd without allocation of unnecessary intermediaries (there is no “save for backward”), fusion variants chosen by measured time, device placement, and plan caching across structurally identical iterations. Outside a tape, optimization is lighter-weight (greedy fusion only), which is basically the “numpy” mode. But you still get to run on all the devices and you get minimalistic fusion too, so for example matmul + relu will be just one kernel, as it should be.

Building from Source

Users installing zyx as a dependency do not need this chapter — cargo add zyx or pip install zyx-py is enough. Prerequisites for every backend are listed in INSTALL.md.

Prerequisites

• Rust 1.88 or newer
• a C compiler (clang or gcc) for the C backend
• optionally, a GPU runtime (CUDA, OpenCL, Vulkan, WGPU); a missing runtime only means that backend contributes no devices
• for the optional Tenstorrent backend, the libttsim_bh.so simulator, the riscv32im-unknown-none-elf target, and llvm-tools

The Tenstorrent firmware lives in zyx/firmware/tt-firmware and is built by zyx/build.rs when the tenstorrent feature is enabled; the ISA tables come from the tt-isa crate. No sibling checkout is required:

git clone https://github.com/zk4x/zyx
cd zyx/zyx
cargo build

The Tenstorrent backend additionally needs the RISC-V target and LLVM tools that build.rs uses for the firmware images:

rustup target add riscv32im-unknown-none-elf
rustup component add llvm-tools
cargo build --features tenstorrent

Tests

No GPU is needed for the CPU backend:

cd zyx
AGENT=1 ZYX_DEV=cpu cargo test --lib --tests -- --skip mnist
AGENT=1 cargo test --doc

AGENT=1 strips ANSI colors; ZYX_DEV=cpu pins the C backend. --skip mnist leaves out the integration test that needs the MNIST dataset in examples/data/. A single test binary is cargo test --test 01_unary, and a single test is cargo test --test 01_unary abs_1 -- --exact.

The Local Check Gate

The full local check gate lives in run_ci.sh at the repository root:

./run_ci.sh

It runs the unit and integration tests, the doctests for every README-bearing crate, rustdoc with private items (so stale doc links fail), the examples, the Python bindings when a torch-enabled interpreter is available, cargo fmt --check, and the book’s code examples.

Examples

cd examples
./tiny-net/run.sh          # simplest
./readme-test/run.sh       # transformer block

Each example has a run.sh that downloads its data if needed and runs in release mode.

The Book Itself

cd zyx-book
./test_book.sh             # dependencies + mdbook test
./test_book.sh --docs      # also rebuild the committed docs/ output

Architecture Overview

zyx has two intermediate representations:

• a tensor graph — an e-graph of equivalence classes, built lazily inside a Tape, where fusion variants compete;
• a kernel IR — a linear SSA list of ops, the unit that passes optimize, autotune searches, and backends compile.

Eager code produces kernel IR directly. Graph code goes through the kernelizer, which turns structural graph nodes into JIT kernels; those kernels are compiled and measured, and extraction then chooses the cheapest measured path through the e-graph.

eager:  Tensor op ──► append to producing kernel ──► compile kernel ──► launch
                                                    (linearize, passes,
                                                     autotune, codegen)

Tape:   Tensor op ──► e-graph ──► kernelizer ──► JIT kernels ──► autotune
                                                                    │
                                              measured times ───────┘
                                                     │
                                            extract cheapest path
                                                     │
                                                Plan (Cmd queue) ──► replay

Both paths end in the same place: a Plan, which is a scheduled command queue (launches, copies, aliases) replayed against memory placements.

Pipeline Stages

1. Tensor API

The user creates tensors and applies operations:

extern crate zyx;
use zyx::{DType, Tensor, ZyxError};
fn main() -> Result<(), ZyxError> {
let x = Tensor::randn([1024, 1024], DType::F32)?;
let y = x.relu();
let z = y * 2.0;
Ok(())
}

A Tensor is a u32 handle into the runtime slab. Outside a Tape, each op is appended to the kernel that produced its inputs; that kernel is compiled and launched when its value is first needed, not at every op. Inside a Tape, ops build graph nodes instead, and Tape::realize compiles and executes the subgraph.

2. The Graph

The graph is an e-graph: Graph { hashcons, ops, jit_kernels, leaf_classes, leaf_map, ... }. Every node belongs to an equivalence class (class_of), and a class chains its alternative forms through next_in_class. Structural hashconsing makes identical expressions share one class.

The graph and the kernel IR share one Op enum (kernel/ops.rs). The graph-level subset is the structural, device-agnostic vocabulary:

Kernel-only variants (Load, Store, Storage, Range, Loop, Wmma, TT, …) are described in The Kernel IR. Shape expressions are interned separately (Runtime::exprs / expr_hash), so symbolic dims are cheap to share and compare.

3. The Kernelizer

The kernelizer (graph/kernelizer.rs) fuses structural graph nodes into kernels. It walks classes in topological order with reference counts: a class is fused into the kernel of its consumer when the surrounding ops allow it, and a value with several consumers is produced once with one output copy per consumer. Storing a value and reloading it for later consumers creates a load kernel at that boundary. fill_gaps runs the same machinery per connected region so that AOT kernels (e.g. a cuDNN matmul) and custom kernels can compete with fused kernels inside the same graph.

Movement ops are not materialized just to be read back. Their index arithmetic is replayed into the consumer, and the linearize pass unfolds it, so a reshape/expand/permute chain costs index arithmetic rather than extra kernels.

The kernelizer’s invariant: after it runs, every output class has a path to the leaves composed only of Program and ToDevice nodes.

4. Autotune Each Kernel

Every JIT kernel produced by the kernelizer goes through the same compilation path as an eager kernel:

1. linearize
2. readability passes (CSE, DCE, instruction scheduling)
3. global-index merge to the device’s limit, renumbering
4. verify
5. three rounds of default_epilogue
6. BeamSearch over default_optimizations, ranked by cost and decided by measured launches
7. device_epilogue (lowering that must not be searched, e.g. Pow expansion and the Tenstorrent physical IR)
8. codegen

The measured time of the winning kernel is recorded on the graph as an Op::Program { info: (ProgramId, nanos) } node. Its program id stays alive; the extraction decides whether to use it.

5. Extract

Graph::extract walks the classes in extraction order and picks, for each class, the producer with the cheapest measured time. There is no estimated cost model: a kernel competes only after it has been compiled and timed. Leaf classes are already realized and need no producer.

After extraction, add_memory_ops inserts ToDevice transfers for the selected producer/consumer pairs that live on different devices.

6. The Plan

Extracted nodes are lowered to a CmdQueue:

• Cmd::Launch — a compiled program plus its argument and output slots;
• Cmd::Copy — a value moved to a destination pool (and device);
• Cmd::Alias — an in-place assign output aliasing its base buffer.

CmdQueue::schedule(outputs) slices the queue into maximal same-device launch runs; copies execute on the destination device, and aliases join the surrounding run as zero-cost bindings. The result is a Plan, cached in Runtime::plan_cache and replayed with the leaf values bound at call time. PlanDim keeps output sizes symbolic until replay, so a plan compiled for one iteration serves later iterations with different variable dims.

7. Backend Codegen

Each backend compiles the winning kernel IR: deSSA, then one linear pass over the ops emitting target code. The IR is deliberately shaped so no backend repeats optimization work. Hardware whose compute model cannot be expressed as a linear SSA pass (Tenstorrent) gets an explicit physical IR below the kernel IR — see Codegen and the Physical IR.

Backends are dispatched by enum, with no dyn Backend:

pub enum Dev {
    Auto,
    C,
    Host,
    Cuda(u16),
    TT(u16),     // feature = "tenstorrent"
    Vulkan(u16),
    OpenCL(u16),
    WGPU(u16),   // feature = "wgpu"
}

Except for the feature-gated WGPU and Tenstorrent backends, every backend is compiled in unconditionally and selected at runtime. A backend whose driver, hardware, or (for Tenstorrent) simulator is missing simply contributes no devices; initialization failures are errors, not panics.

8. Runtime

The runtime (runtime.rs) is the process-wide state: graphs, tensors, kernels, programs, timings, interned shape expressions, RNG, and the plan cache. A global Mutex<Runtime> guards it.

Device selection for generic work is Runtime::pick_device(bytes): start from Dev::all(), drop AOT-only devices, keep devices whose pool has at least bytes free, and pick the one with the most free compute capacity. Cross-device movement is scheduled as copies in the plan, and asynchronous execution is a per-backend concern (CUDA’s worker micro-batches launches) — see Runtime and Plans.

Debugging the Pipeline

ZYX_DEBUG=4 and ZYX_DEBUG=8 are complementary: the first shows the pre-linearize DAG (structural ops, vector values), the second the post-search linear form the backend receives. ZYX_TT_IR=1 dumps the Tenstorrent physical IR phases.

Key Design Decisions

• One graph for everything — autograd and computation share the same graph. No need to specify which tensors require gradients.
• Symbolic dims everywhere — the eager path and the graph path both work with symbolic dimensions, always. Every symbolic dim bottoms out in a Param { Variable } scalar (IDX_T) whose value lives in a backend pool’s variable slot. Consumers evaluate dims by folding Const leaves and reading variable slots, never by fabricating placeholders (0, -1, …).
• Measured search, no cost model — kernel variants are ranked by a cheap feature cost, then decided by real launches; extraction also uses measured times. This keeps rewrites exact: the winner is the one that actually ran fastest.
• Inline ops — all kernel ops live in the arena as flat entries: Op is 24 bytes, OpLinked 32. No Box, no vtables, no per-op indirection.
• Linear IR — a linked list of fixed-size ops; passes traverse it front-to-back or back-to-back.
• Backend codegen is trivial — the hard work lives in the kernel IR search and passes.
• Tape-scoped lazy graph — ops inside a tape build graph nodes instead of executing eagerly. This enables fusion search, device placement, and plan caching across iterations.

The Tensor

The Tensor is the primary user-facing type in zyx. It is designed to be a lightweight handle — only 4 bytes.

pub struct Tensor {
    pub(super) id: TensorId,  // u32 index into the runtime's tensor slab
}

Design Choices

Why 4 Bytes?

Most ML frameworks have heavyweight tensor objects. PyTorch’s Tensor is a TensorImpl* with shape, stride, dtype, device, storage, and autograd metadata — easily 100+ bytes. In zyx, metadata lives in the runtime’s slabs — tensor state, interned shape expressions, and kernel/program tables — not in the handle. The tensor is just an index.

Reference Counting

Tensors are reference-counted via the global RT (a Mutex<Runtime>):

impl Clone for Tensor {
    fn clone(&self) -> Self {
        RT.lock().retain(self.id);
        Tensor { id: self.id }
    }
}

If we used Arc instead, we would still need Mutex for the Runtime — Tensor(id, Arc<Mutex<Runtime>>). The current approach avoids the Arc overhead and keeps Tensor at 4 bytes. Since every tensor operation already locks the runtime to append a graph op, there’s no additional lock contention from reference counting.

Execution

Outside a tape, each op is appended directly to the kernel that produced its inputs. When fusion is not possible, the kernel compiles and executes:

extern crate zyx;
use zyx::{DType, Tensor, ZyxError};
fn main() -> Result<(), ZyxError> {
let x = Tensor::randn([1024, 1024], DType::F32)?;
let y = x.relu();     // appended to x's kernel
let z = y.tanh();     // appended to same kernel
// at some point the kernel compiles and executes
Ok(())
}

Inside a tape, operations build graph ops lazily and execute when the tape is realized (dropping only cleans up graph state — no computation is performed). The key insight: repeated graph patterns are automatically recognized and cached across structurally identical iterations.

Construction Methods

Tensors can be created from:

extern crate zyx;
use zyx::{DType, Tensor, ZyxError};
fn main() -> Result<(), ZyxError> {
let t = Tensor::from([1.0f32, 2.0, 3.0]);
let t = Tensor::randn([1024, 1024], DType::F32)?;
let t = Tensor::uniform([1024, 1024], -1.0f32..1.0)?;
let ones = Tensor::ones([3, 3], DType::F32);
let zeros = Tensor::zeros([3, 3], DType::F32);
Ok(())
}

This also works from files on disk: Tensor::load_safetensors(path)? returns a map of tensors backed by the disk pool, read only when a plan needs them.

Values, not Buffers

Every op returns a new tensor; nothing mutates behind your back. The one explicit in-place write is Tensor::assign, which records a store in the graph:

extern crate zyx;
use zyx::{DType, Tensor, ZyxError};
fn main() -> Result<(), ZyxError> {
let x = Tensor::randn([3, 3], DType::F32)?;
let x_plus_one = &x + 1.0;  // new tensor, no mutation
Ok(())
}

Autograd never has to track hidden mutation: an assign is an op in the graph like any other, so the backward pass sees it.

The Graph

The graph is an e-graph (equivalence graph) used in tape mode for fusion, autograd, and device placement. Inside a Tape, every operation builds an Op in this graph. Outside a tape, there is no graph — ops go directly to kernel fusion. When a tape is active, the graph is shared between computation and autograd — there is only one.

Data Structure

The graph is stored in the Runtime:

pub struct Graph {
    pub(crate) hashcons: Map<Op, OpId>,
    pub(crate) ops: Slab<OpId, OpNode>,
    pub(crate) jit_kernels: Slab<JitKernelId, JitKernelData>,
    pub(crate) leaf_classes: Vec<OpId>,
    pub(crate) leaf_map: Map<OpId, TensorId>,
    pub(crate) ref_count: u64,
    pub(crate) dead: bool,
    pub(crate) max_cons_id: u32,
}

pub(crate) struct OpNode {
    pub(crate) op: Op,
    pub(crate) class_of: OpId,
    pub(crate) next_in_class: OpId,
}

The hashcons deduplicates structurally identical ops — if the same operation on the same inputs already exists, the existing OpId is reused. This provides CSE (common subexpression elimination) for free.

Each OpId maps to an OpNode entry in the ops slab. A node joins exactly one equivalence class, named by its class_of field (the class’s root OpId); alternative forms of the same value are chained through next_in_class, in insertion order. max_cons_id allocates the identity used by Op::Param::id.

Op Types

The graph opset is derived from tinygrad. By stacking these variants, zyx can express ALL linear algebra operations and ALL PyTorch ops:

enum Op {
    Const(Constant),
    Param { dtype: DType, kind: ParamKind, shape: OpId, id: u32 },
    Reshape { x: OpId, shape: OpId },
    Expand { x: OpId, shape: OpId },
    Permute { x: OpId, axes: TinyVec<UAxis> },
    Pad { x: OpId, axis: UAxis, lp: OpId, len: OpId },
    Flip { x: OpId, axes: TinyVec<UAxis> },
    Narrow { x: OpId, axis: UAxis, start: OpId, len: OpId },
    Stack { ops: Box<[OpId]> },
    Index { vec: OpId, idx: usize },
    Reduce { x: OpId, rop: BOp, reduce_axis: OpId },
    Cast { x: OpId, dtype: DType },
    Bitcast { x: OpId, dtype: DType },
    Unary { x: OpId, uop: UOp },
    Binary { x: OpId, y: OpId, bop: BOp },
    Select { cond: OpId, if_true: OpId, if_false: OpId },
    Store { dst: OpId, dst_idx: OpId, src: OpId },
    After { x: OpId, dep: OpId },
    ToDevice { x: OpId, device: Dev, time: u64 },
    Contiguous { x: OpId },
    Detach { x: OpId },
    Program { inputs: OpId, outputs: OpId, info: Box<(ProgramId, u64)> },
    Kernel(Box<CustomKernel>),
    // Kernel-only variants (Load, Storage, Range, Loop, Wmma, TT, ...)
    // are described in The Kernel IR.
}

All inputs reference OpId — ops operate on equivalence classes. View ops (Reshape, Expand, Permute, Pad, Flip, Narrow) are per-axis and stackable; Stack builds vectors; Index selects an element; Bitcast reinterprets bits without value conversion; Store is an in-place write (assigns lower to it); After orders side-effecting ops; Contiguous forces a materialization boundary; Detach is value identity with no gradient; Program is a compiled-kernel boundary; Kernel wraps an opaque user custom kernel.

The reason to use OpId instead of TensorId is to get canonicalization for free.

Lifecycle with Tape

There is no graph outside a tape — ops are fused directly into kernels.

Inside a tape, ops accumulate until Tape::realize() or drop. A class can hold several alternative producers, and extraction chooses between them:

• CSE via hashconsing: identical structural ops share one class.
• Kernel twins: the kernelizer mints an Op::Program node for every fused kernel candidate, on the class it produces.
• AOT pattern matches: backends add Op::Program nodes for precompiled kernels (cuDNN, …), so they compete with fused kernels.
• Transfers: extraction inserts Op::ToDevice nodes where the chosen producer and consumer live on different devices.

There is no cost model: each candidate is compiled and autotuned, and Graph::extract picks the path by measured timing. Realized ops that the tape references are preserved for autograd; unreferenced ops are released.

Graph Size

The graph is designed to stay small. Tensor handles are u32 (4 bytes), so a vector of 10,000 handles costs ~40 kB. The graph nodes themselves are bigger: Op is 24 bytes and an OpNode is 32. When the tape is dropped, the graph is deleted as a whole, so there is no per-op garbage collection. Per-tensor rc and the graph’s ref_count decide when that happens; ops are otherwise append-only (arena style).

The Kernelizer

The kernelizer (graph/kernelizer.rs) fuses tensor operations into kernels. In eager mode, ops are appended directly to kernels. In tape mode, it traverses the accumulated graph and fuses compatible nodes into kernels.

When the Kernelizer Runs

Outside a tape, there is no graph kernelizer: Runtime appends each op to the kernel that produced its inputs and splits (duplicate_or_store) only when a fusion boundary is hit. Inside a tape, Graph::kernelize runs once during Tape::realize() (dropping only cleans up graph state):

extern crate zyx;
use zyx::{DType, Tape, Tensor, ZyxError};
fn main() -> Result<(), ZyxError> {
// Tape mode — kernelizer runs at realize
let x = Tensor::randn([4], DType::F32)?;
let tape = Tape::new([&x])?;
let y = x.relu();
let z = y * 2.0;
// at this point: no computation done yet
tape.realize([&z])?;
// kernelizer runs → optimizes → executes
Ok(())
}

Fusion Logic

In tape mode, the kernelizer processes graph ops bottom-up in topological order. Each graph op type adds ops to the kernel its input lives in — fusion is the default, splitting happens only when needed.

Per-Op Type Behavior

When Splitting Happens

The key splitting decision is in duplicate_or_store_class(). When a kernel has multiple outputs (a tensor used by >1 downstream), the kernelizer chooses:

1. If preceded by a reduce (expensive to recompute) → store the intermediate result to global memory, create a new load kernel for the next consumer
2. If NOT preceded by a reduce (cheap to recompute) → duplicate the kernel so each consumer gets its own copy

This is a cost heuristic, not a simple rule: is recomputing cheaper than a global memory store+load?

Requested outputs are consumers too: if a class is needed after its producing kernel is closed, it is stored and a fresh load kernel serves the remaining consumers.

Practical Outcome

Each kernel tends to center around one reduce loop, with all fused element-wise ops grouped before and after it. A chain of element-wise ops always fuses into one kernel. Reduce-heavy graphs may end up with each reduce in its own kernel separated by store/load boundaries.

The Kernelizer Struct

The kernelizer is a set of methods on Graph; all working state is local to a call:

impl Graph {
    pub fn kernelize(&mut self, inputs: &Set<OpId>, outputs: &BTreeSet<OpId>, allowed: Option<&Set<OpId>>) {
        // `jit_kernels: Slab<JitKernelId, JitKernelData>` is the output table.
        let mut rcs: Map<OpId, u32> = Map::default();
        let mut visited: Map<OpId, (JitKernelId, OpId)> = Map::default();
        // ...
    }
}

pub struct JitKernelData {
    pub(crate) kernel: Kernel,
    pub(crate) outputs: Vec<OpId>,
    pub(crate) loads: Vec<OpId>,
    pub(crate) stores: Vec<OpId>,
}

The visited map tracks which graph ops have been converted to kernel ops. When a graph op is already in visited, the kernelizer uses the existing kernel result instead of recomputing — this is how shared subgraphs are handled.

Building Kernel Ops

Each graph op type maps to kernel IR operations:

View operations (reshape, expand, permute, pad) are unfolded into index arithmetic rather than becoming separate ops. This is how they become “free” — the index computation is inlined into the load/store operations.

Kernel Caching

Kernels are deduplicated by content, and compiled programs are cached by kernel id:

// Runtime
kernel_map: Map<Kernel, KernelId>,          // content-dedup: equal kernels share an id
programs: Map<KernelId, DeviceProgramId>,   // compiled program per kernel id
timings: Map<ProgramId, u64>,               // measured nanos

When a kernel finishes optimization its hash is computed; an equal Kernel already in kernel_map reuses the KernelId (and the compiled program, if one exists). Compiled programs persist across realizations, so repeated graph patterns — training-loop iterations — hit the cache after the first one. The plan cache (Runtime::plan_cache) does the same for whole execution plans, keyed by graph content plus leaf pools.

The Kernel IR

The kernel IR is the representation that passes optimize and backends compile. It is a linear SSA construct: a doubly-linked list of ops stored in an arena. This chapter describes its structure and its ops; the passes that transform it are in Optimization Passes.

Data Structure

pub struct Kernel {
    pub(crate) ops: Slab<OpId, OpLinked>,   // arena
    pub head: OpId,                         // first op
    pub(crate) tail: OpId,                  // last op
    pub(crate) dev: Dev,                    // target device
    pub(crate) dev_info: Option<Arc<DeviceInfo>>, // fixed hardware facts
}

pub struct OpLinked {
    pub prev: OpId,   // u32
    pub next: OpId,   // u32
    pub op: Op,       // 24 bytes (enum + payload)
}

OpId is a u32 index: random access is O(1), OpId::NULL is u32::MAX, and the list is walked through head/next (or backwards from tail). Op is 24 bytes and OpLinked 32 — both are compile-time assertions in kernel/ops.rs. There is no Box per op and no vtable.

dev_info is a snapshot of the device’s DeviceInfo (warp size, limits, capabilities) taken when the device is bound. Passes never lock the runtime to ask hardware questions, and clones share the snapshot through Arc.

The graph uses the same Op enum, but a different node: a graph class node carries class_of and next_in_class instead of prev/next. The kernel node type is OpLinked; the graph node type is OpNode.

Linearization

Before linearize, a kernel is a small DAG with no loads and no storage:

• Param { kind: Global } — an input buffer;
• Param { kind: GlobalMut } — an output buffer / store destination;
• Param { kind: Variable } — a scalar launch argument (e.g. a dynamic dim);
• movement ops (Reshape, Expand, Permute, Flip, Pad, Narrow);
• element-wise ops (Unary, Binary, Select, Cast, Bitcast, Reduce, …);
• Store at the end.

linearize turns that into the linear form:

• movement and Reduce ops are removed and expanded into index arithmetic;
• Load ops are inserted for read-only params, with Storage as their base;
• Index/arithmetic over group and loop indices computes addresses;
• loops (Loop/EndLoop) and ranges (Range) become explicit;
• every Store gets an index expression computed from the shape it writes;
• writable params stay in place as store destinations.

A real example makes the difference visible. Before linearize (ZYX_DEBUG=4):

r0: i64 = 3
r1: f32[3] = param global shape=r0
r2: f32 = 0
r3: f32[3] = expand r2 -> ...
r4: bool[3] = r1 > r3
r5: f32[3] = f32(r4)
r6: f32[3] = param global shape=r0
r7: f32[3] = r5 * r6
...
r17: f32[1] = pad r14 axis=0 lp=r15 len=r16
r18: f32[1] = param global mut shape=r16
store vec1 r18 = r17

After linearize and the search (ZYX_DEBUG=8), the same computation is scalar ops over explicit indices:

r85: f32 = param global shape=NULL
r18: f32 = param global mut shape=NULL
r52: f32 = 0
r31: i64 = 0
r51: f32 = load r85[r31]
r9: f32 = -r51
r11: bool = r9 > r52
r4: bool = r51 > r52
r13: f32 = r11 ? r9 : r52
r7: f32 = r4 ? r51 : r52
r14: f32 = r7 + r13
store r18[r31] = r14

Two details matter:

• Param shapes go to NULL during linearization. Buffer sizes must be resolved before it (the allocator and launch-argument code do exactly that); after it, a param’s shape is no longer how sizes are communicated.
• The load/store index carries the access shape. load r85[r31] with a scalar index reads one element; an index over a lane vector reads a vector; an index over two distinct axes reads a tile. The debug printer spells this out as store vec4 or store tile 32x32.

Kernel::is_linearized detects the state from the stores: a pre-linearize store writes a whole view to a bare Param; linearization gives every store destination an index expression.

Op Variants

Leaves

Op::Const(Constant)
Op::Param { dtype: DType, kind: ParamKind, shape: OpId, id: u32 }
// ParamKind::Variable  — scalar launch argument (e.g. a dynamic dim)
// ParamKind::Global    — read-only buffer argument
// ParamKind::GlobalMut — read-write buffer argument

Param::id is graph identity for hashconsing: two params name the same buffer iff their ids agree. Kernel passes ignore it — kernel params are bound positionally at launch — and hashing canonicalizes ids to head-order ordinals so program caches keep sharing across buffers.

Arithmetic and selection

Op::Cast { x: OpId, dtype: DType }
Op::Bitcast { x: OpId, dtype: DType }
Op::Unary { x: OpId, uop: UOp }
Op::Binary { x: OpId, y: OpId, bop: BOp }
Op::Select { cond: OpId, if_true: OpId, if_false: OpId }
Op::Stack { ops: Box<[OpId]> }       // pack values into a vector
Op::Index { vec: OpId, idx: usize }  // select one element of a vector

There is no Mad op: Kernel::mad(x, y, z) builds a Binary::Mul followed by a Binary::Add, and the passes may reassociate/fold it. There is no separate “fused multiply-add” IR node.

Memory

Op::Storage { dtype: DType, scope: MemScope, len: Dim, id: u32 }
Op::Load  { src: OpId, src_idx: OpId }
Op::Store { dst: OpId, dst_idx: OpId, src: OpId }
Op::Copy  { src: OpId, src_idx: OpId, dst: OpId, dst_idx: OpId }

Storage is kernel-internal memory — registers, local/shared memory, Tenstorrent circular buffers. Its id is assigned by the register allocator: before allocation every storage of the same (dtype, scope, len) is one request and compares equal; after allocation id separates physical slots so CSE cannot merge two registers that must stay distinct.

Load is memory → SSA. Store is SSA → memory and never memory → memory; that is Copy. Store and Copy are effect-only: they produce no SSA value, are DCE roots, and are never CSE’d or hoisted.

Iteration

Op::Range { axis: u32, kind: RangeKind }
Op::Loop { len: OpId }
Op::EndLoop
Op::Barrier
Op::Unroll { axis: u32, len: u16 }

RangeKind has two variants:

• Group(OpId) — a grid axis (CUDA blocks, CPU cores, Tenstorrent cores);
• Local(u32) — a thread/lane axis. Local axis 0 is the hardware warp; there is no separate warp kind, and the lane id is the axis-0 local value.

There is no Scalar range kind: a scalar is a rank-0 value, not a range.

Loop is a counted loop. A conditional is a Loop with a boolean length (0 or 1 trips), closed by EndLoop exactly like a counted loop; there is no If/EndIf. Barrier pins program order as an operand-free effect and DCE root. Unroll is a lane vector produced by splitting: it holds the lane constants 0..len tagged with the axis they were split from, and is devectorized into a Stack of constants right before codegen.

Compute

Op::Wmma { dims: MMADims, layout: MMALayout, dtype: MMADType, a: OpId, b: OpId, c: OpId }
Op::Reduce { x: OpId, rop: BOp, reduce_axis: OpId }   // pre-linearize, or over lanes
Op::Asm { asm: TinyString, ops: TinyVec<OpId> }

Wmma is the fused fragment matmul d = a @ b + c (the result is the new accumulator). Reduce is a pre-linearize op: linearize rewrites it into a loop with an accumulator. It reappears after split turns a loop into lane vectors, where the reduction runs over the lanes rather than over a loop; that is why the variant stays in the enum. Asm emits a backend assembly template with {i} operand substitution.

View and movement ops

Op::Reshape { x: OpId, shape: OpId }
Op::Expand  { x: OpId, shape: OpId }
Op::Permute { x: OpId, axes: TinyVec<UAxis> }
Op::Flip    { x: OpId, axes: TinyVec<UAxis> }
Op::Pad     { x: OpId, axis: UAxis, lp: OpId, len: OpId }
Op::Narrow  { x: OpId, axis: UAxis, start: OpId, len: OpId }

These are pre-linearize ops. They are replayed into consumers and unfolded into index arithmetic by linearize.

Graph-only markers

Op::After { x: OpId, dep: OpId }          // x may not run before dep
Op::ToDevice { x: OpId, device: Dev, time: u64 }
Op::Contiguous { x: OpId }                // fusion-break hint
Op::Detach { x: OpId }                    // autograd break; erased by kernelizer

After and ToDevice survive extraction; Contiguous forces an output boundary; Detach is value identity and the kernelizer remaps it away.

Opaque kernels and backend payloads

Op::Source(SourceBlock)                        // program source block
Op::Program { inputs: OpId, outputs: OpId, info: Box<(ProgramId, u64)> }
Op::Kernel(Box<CustomKernel>)                  // user custom kernel boundary
Op::TT(TTOp)                                   // Tenstorrent physical IR
Op::GPU(Box<GPUOp>)                            // shared GPU launch config
Op::Spirv(Box<SpirvOp>)                        // Vulkan SPIR-V payload
Op::PTX(Box<PTXOp>)                            // CUDA direct-PTX payload

Program’s inputs and outputs are Stack classes, not per-kernel boxes; info carries the program id and its measured time. Timings are ignored by Eq/Hash.

Memory Scopes

pub enum MemScope {
    Global,    // DRAM / device global memory
    Local,     // SRAM / shared memory
    Register,  // registers
    Circular,  // circular buffer, SRAM (Tenstorrent)
}

There is no Variable scope: scalar variables are Param { kind: Variable }, not storage.

Shapes

Kernel::shape(op) reports the value shape of an op as a rank list (at most eight dims): Unroll contributes one dim of len lanes on its axis, and a Param contributes its buffer shape’s dims; everything else is scalar. For loads, stores, and copies, the index’s shape is the access shape, which is what the debug printer shows (store vec1, store tile 32x32).

Loops and ranges are scalar iterators and contribute no element dims: the trip count is loop structure, not value shape.

Codegen

Codegen consumes a stabilized kernel:

1. deSSA — resolve SSA values to physical registers/memory slots;
2. linear pass — walk the op list once, emitting instructions.

No further optimization happens here. Backends whose hardware needs explicit physical state get one more IR below this one — for Tenstorrent, the tt_tile → tt_dataflow_select → SFPU/LREG/DST passes → tt_verify sequence in device_epilogue, ending in Op::TT ops. See Codegen and the Physical IR.

Debugging

ZYX_DEBUG=4 cargo run   # kernel IR before linearization
ZYX_DEBUG=8 cargo run   # winning kernel IR after the search
ZYX_DEBUG=16 cargo run  # generated source/assembly

ZYX_TT_IR=1 additionally dumps the Tenstorrent physical IR before and after tt_tile, tt_reindex, and tt_dataflow_select. Debug output is colorized only on an interactive terminal; AGENT=1 disables color.

Kernel::verify checks ordering, dtypes, shapes, and use/def consistency. It is compiled in but returns immediately outside debug_assertions. After the Tenstorrent phase boundary the IR is native tiles/rows, and tt_verify is the checker instead.

Optimization Passes

Optimization passes operate on the linear kernel IR. There are two kinds:

• always-on passes — the epilogue, run on every candidate and before compilation;
• autotuned passes — made into a config space by a MakeOpt and selected by BeamSearch using measured launches.

The design goal of the small opset is that new passes are easy to write. There is no fixed number zyx aims to have; passes are added when a performance opportunity is identified.

The Epilogue

default_epilogue is the always-on pipeline. It runs after each optimization step during the search and three times before the search starts:

pub fn default_epilogue(&mut self) {
    self.delete_useless_loops();
    self.constant_folding();
    self.move_constants_to_beginning();
    self.loop_invariant_code_motion();
    self.fold_accs();
    self.algebraic_simplifications();
    self.simplify_accumulating_loop();
    self.simplify_accumulating_loop();
    // Second sweep: the first sweep's rewrites expose more work.
    self.constant_folding();
    self.move_constants_to_beginning();
    self.loop_invariant_code_motion();
    self.fold_accs();
    self.algebraic_simplifications();
    self.simplify_accumulating_loop();
    self.swap_commutative();
    self.common_subexpression_elimination();
    self.instruction_schedule();
    self.dead_code_elimination();
}

The core runs twice because a first pass of constant folding/LICM/algebraic simplification often exposes a second round of the same rewrites. simplify_accumulating_loop appears twice in the first sweep for the same reason. The epilogue is device-independent: it never converts a transcendental and leaves Pow alone.

What each always-on pass does

• Delete useless loops — collapses length-1 counted loops and ranges to the constant index 0 (the body still executes once), then removes loops whose body has no effect. Boolean-length loops are conditionals, not trip counts, and are never collapsed by the length-1 rule.
• Constant folding — evaluates any op whose inputs are all constants (2 + 3 → 5, sin(0.0) → 0.0).
• Move constants to beginning — hoists constant definitions to the head, which exposes fusion opportunities and simplifies register allocation.
• Loop-invariant code motion (LICM) — hoists ops that produce the same value on every iteration.
• Fold accumulators — removes accumulator updates whose value is never stored back inside the loop (dead accumulation).
• Algebraic simplifications — x * 1 → x, x + 0 → x, x - x → 0, and similar identities over the concrete dtype.
• Simplify accumulating loop — replaces a closed-form accumulate loop (for example a sum of successive indices) with its arithmetic formula; it processes one such loop per call.
• Swap commutative — canonicalizes commutative operands so CSE hashes see identical forms.
• Common subexpression elimination (CSE) — merges structurally identical ops. Storage allocation ids and param ids are part of the key, so two physical registers or two buffer arguments are never merged after allocation.
• Instruction scheduling — orders the head as read-only params, GlobalMut params, local read-only storages, then local read-write storages; the rest is topologically sorted with fewest dependencies first. Memory ops that share a base keep their relative order, stores never leave their loop/conditional, and nothing moves before a barrier.
• Dead code elimination (DCE) — removes ops with no users. Store, Copy, Barrier, Loop, and program-order roots are always kept. DCE is last: backends expect no unreferenced values.

The Device Epilogue

device_epilogue runs exactly once per kernel, after the search and before the timed launches and codegen. These decisions are lowering, not search: a kernel optimized ten times during the search must not be re-expanded each time, and a user-supplied epilogue must not be able to skip them.

pub fn device_epilogue(&mut self) {
    self.unfold_pows();               // Pow -> exp2(y * log2 x), every backend
    if device_is_tenstorrent {
        self.tt_tile();               // vector IR -> tiles
        self.tt_unfold();             // Exp2/Log2 become native ops
        self.tt_reindex();            // logical indices -> tile traffic
        self.default_epilogue();      // clean the tile IR
        self.tt_dataflow_select();    // phase boundary: native instructions
        self.tt_sfpu_fusion();        // value -> lane register
        self.tt_lreg_alloc();         // virtual -> physical LREGs
        self.tt_dst_reuse();          // coalesce DST tiles
        self.tt_reg_sync();           // Semset/Wait/Post hand-offs
        self.tt_verify();
    } else {
        self.devectorize();           // lane vectors -> scalar statements
    }
}

The Tenstorrent path is the physical IR described in Codegen and the Physical IR; generic passes may not run after tt_dataflow_select, and tt_verify replaces the generic verify.

Autotuned Passes

default_optimizations returns the make functions the search explores:

pub const fn default_optimizations() -> [MakeOpt; 10] {
    [
        Kernel::opt_split,
        Kernel::opt_reassociate_commutative,
        Kernel::opt_local_reduce,
        Kernel::opt_merge,
        Kernel::opt_pad,
        Kernel::opt_tensorize,
        Kernel::opt_local_staging,
        Kernel::opt_recipe,
        Kernel::opt_register_tiling,
        Kernel::opt_tile_and_stage,
    ]
}

A make function scans the kernel at a stable state and returns a Box<dyn Optimization>: an object with nconfigs() and apply(kernel, config). No pass runs between make and apply, so OpIds embedded by the make function stay valid. An optimization that cannot fire returns zero configs or bails inside apply; it never aliases a wrong transformation.

Split

opt_split splits one site per config: a group range into a smaller group plus local threads (SplitMode::Local), a group range into a group plus unroll lanes (SplitMode::Unroll), or a loop into two nested loops (SplitMode::Loop). The split index is composed as an affine expression inline, so the IR grows only in op count, not in vector width; lane growth happens later when devectorization desugars the chosen variant. Validity — divisibility, lane limits, the boolean-loop gate — is checked at make time.

Reassociate commutative

opt_reassociate_commutative re-associates commutative chains so that constants and repeated operands group together. It exposes folding and CSE opportunities that operand order alone hides.

Local reduce

opt_local_reduce targets a single large loop reducing into one register accumulator (a standalone sum/max over a long axis, not a matmul). It parallelizes the reduction across threads, stores partial results in local memory, and tree-reduces them. One config per (reduce site, local factor, global factor).

Merge

opt_merge fuses exactly two axes into one flat axis and rewrites the originals as div/mod decomposition of the flat axis. Each config merges one nested loop chain (inner loop inside the outer’s region, no store/copy in between) or one range pair (both group or both local; axes ascending; symbolic lengths bail). Chains fold through pairwise calls, so a refused later step leaves valid IR.

Pad

opt_pad rounds range lengths up to a hardware-friendly multiple (for example 32) using the round-up idiom as IR, with masks so out-of-range loads read zero and out-of-range stores are skipped. Loops pad only on Tenstorrent, where tile loops must run full-tile trip counts. Lengths already provably a multiple are left untouched, and padding is idempotent.

Tensorize

opt_tensorize recognizes a scalar matmul accumulation loop (acc = a[m,k] * b[k,n] + acc over K) and rewrites it into Op::Wmma. It emits one config per device MMA variant the extent and dtypes fit (DeviceInfo::mma_variants). The pass does the tiling and fragment addressing itself; the result is scalar IR of slot-wise register accesses, so devectorize has nothing to undo, and it removes the original scalar region.

Local staging

opt_local_staging deduplicates loads that threads repeat. When a load’s index does not mention a local axis, every thread re-loads the same value; staging fills a MemScope::Local tile cooperatively once and lets the threads read it back. The pass strip-mines the innermost loop the load depends on and splits the body into a fill half and a compute half (ops between the load and its consumers that depend only on the load and constants move into the fill, so a dequantized value is staged rather than its packed form). One config per (repeated load, strip factor); it bails on devices without local memory, non-affine indices, or tiles that do not fit or divide over the block.

The recipes

The last three passes are recipes: fixed sequences of known-good steps exposed as one search move, so the beam can reach a good kernel shape without walking every intermediate state. A recipe returns a config space, never a winner, and only fires from a plain kernel (no lanes, barriers, reduces, local memory, or fused matmul — another pass already ran).

• opt_recipe — the lane recipe: split a group range into local lanes (factors tried largest first), then stage the operand the lanes share.
• opt_register_tiling — the register-blocking recipe (the former coarsen era): unroll the reduce loop into accumulators, then each output range into a per-thread tile. Reduce sites are mandatory; output ranges may be skipped.
• opt_tile_and_stage — lanes on one output range, a per-thread tile on both, then stage the shared operand.

How the Search Works

BeamSearch is the driver. Its defaults are:

BeamSearch {
    n_launches: 1,          // candidates actually launched and timed
    n_seeds: 200,           // initial optimization samples
    n_added_per_step: 200,  // new candidates per expansion
    n_removed_per_step: 0,  // candidates dropped per step
    n_total_opts: 1000,     // expansion iterations
}

The flow:

1. The caller supplies seeds: linearized kernels, already epilogue-cleaned, all on one device and sharing one parameter signature. The search does no seed preprocessing.
2. Every make function is applied to the base kernel; configs are subsampled proportionally to n_seeds, the epilogue runs after each apply, and the resulting kernel hash deduplicates visited states.
3. Candidates are ranked by cost (by default Kernel::base_cost, a weighted feature count — no launch). The best-scoring sequence is sampled and expanded with new configs; n_total_opts expansions are attempted, dropping the worst n_removed_per_step each round.
4. The top n_launches candidates are rebuilt deterministically, device_epilogue runs once on each, and each is compiled and launched with launch_with_timings. Candidates that fail to compile/launch are skipped, and the search keeps going until it has measured n_launches successes or runs out.
5. Every measured program is released; the fastest kernel is returned with its time. The cost function only orders the search — measured time decides.

Kernels are identified by content hash (Kernel::get_hash), so the same state reached by different sequences is explored once. ZYX_DEBUG=256 selects NoSearch, which skips the seed search and compiles a cleaned seed once — useful to bisect whether a wrong result comes from the search space or from seed preparation/codegen.

Performance Budget

The repeated passes must stay fast: 30 microseconds per invocation is the ceiling for default_epilogue and for the whole default_optimizations list per candidate. Those run once per candidate inside the search, so a slow pass compounds into seconds of autotuning. The device epilogue runs once after the search and is not bound by that budget.

Passes are measured with the time_pass! macro or std::time::Instant on stderr; a regression in algorithmic cost is a performance bug, not a style issue.

Correctness Guarantee

No optimization is needed for tests to pass. All tests must pass no matter which sequence of optimizations (including empty) is applied.

If an optimization breaks a test, the optimization is buggy — not the sequence ordering. Passes must preserve semantics or bail (todo!()/zero configs), never silently alias a wrong transformation. Kernel::verify checks IR consistency in debug_assertions builds; the Tenstorrent physical IR is checked by tt_verify.

Backend System

Zyx supports multiple hardware backends. Backends are enum-dispatched, compiled into the library (WGPU and Tenstorrent are feature-gated), and selected at runtime through a lightweight Copy handle.

Dev — Device Handle and Selector

There is no dyn Backend and no device ids separate from the devices: Dev is both the selector and the handle — a Copy enum naming the backend plus the hardware ordinal.

pub enum Dev {
    /// Auto-select: resolves to the first available device (Dev::all).
    Auto,
    /// CPU backend (runs on Pool::Host).
    C,
    /// Transfer-only pseudo-device (runs on Pool::Host); never enumerated.
    Host,
    /// CUDA GPU with the given driver ordinal.
    Cuda(u16),
    /// Tenstorrent chip with the given id.
    TT(u16),          // feature = "tenstorrent"
    /// Vulkan physical device with the given index.
    Vulkan(u16),
    /// OpenCL device with the given index.
    OpenCL(u16),
    /// WGPU device with the given index.
    WGPU(u16),        // feature = "wgpu"
}

Trait objects would require downcasting to reach backend-specific functionality; the enum makes every method an exhaustive match instead. Dev::Auto is the scheduling placeholder — it is resolved by the scheduler before placement, so the device API never sees it. The panic! arms for Auto in methods like info() and free_compute() are guards on that invariant, not behavior; and the device’s memory pool is always derived from the device via Dev::pool() — never the reverse.

Pool — Memory

Memory belongs to pools, and pools are process-wide: they outlive any Runtime and are never deinitialized.

pub enum Pool {
    /// Host RAM. Shared by the C and Host devices, which own no pool.
    Host,
    /// Disk-backed tensors (paths, not bytes).
    Disk,
    Cuda(u16),
    OpenCL(u16),
    Vulkan(u16),
    TT(u16),          // feature = "tenstorrent"
    WGPU(u16),        // feature = "wgpu"
}

Each variant owns its globals — one Mutex<pool> per ordinal (Host is a OnceLock<Mutex<_>>, Disk is a const Mutex with no OnceLock at all, the rest are OnceLock<Vec<Mutex<_>>> with one Mutex<()> per backend serializing first construction where it matters — cheap backends have no init lock). All globals live at the top of each backend/*.rs file, so the process-wide state is visible before any logic. The only lock takers are the device-API entry points (alloc/free/copy/compile/launch); the per-op tensor path never touches them.

Lazy Initialization

There is no upfront backend-initialization phase. Dev::all() triggers lazy init of every backend; backends that are configured out, whose driver is missing, or whose hardware is absent contribute nothing to the returned list. This keeps startup free and makes device discovery idempotent. By default, zyx will search for all devices and use the fastest for the computations. If only one device is accessed by manually using Tensor::to(Dev), then only that device is initialized and other devices are left uninitiaized. So initialization is lazy and happens at the latest possible moment.

impl Dev {
    pub fn all() -> Vec<Dev> {
        // C: single device, included if init succeeds.
        // CUDA / TT / Vulkan / OpenCL / WGPU: one entry per detected device.
    }
}

Device selection happens at schedule time: Auto resolves to the first available device; a kernel is placed on a device whose pool has enough free bytes, skipping AOT-only devices (see below) for generic kernels.

Device API

Every backend implements the same surface, dispatched by the enum:

pub fn compile(self, kernel: &Kernel, debug_asm: bool) -> Result<DeviceProgramId, BackendError>;
pub fn launch_timed(self, program_id: DeviceProgramId, args: &[LaunchArg]) -> Result<u64, BackendError>;

plus pool operations (alloc, free, release, pool_to_host, copy) and info(). One backend-specific concept lives in the shared API: GwsDim — the per-axis global work size. Each gws dim is a Group index length, either a constant or a Param-backed dynamic length resolved from launch args; how it maps to a launch grid (CUDA grid, OpenCL global size, …) is each backend’s own business, derived from its Op::Range ops at compile time.

Codegen: Mostly Trivial, by Construction

For most backends, codegen is a straight line: deSSA, then one linear pass over the kernel IR emitting target code. This is not an accident — the kernel IR is optimized until nothing searchable remains (see the search test): everything the autotuner could enumerate over has already been decided, measured, and frozen before codegen runs.

What is left for a backend is exactly the non-searchable, non-SSA residue: registers, sync and placement concerns, launch-argument conventions, format configs — machine-shaped facts the value-level IR cannot and should not express. For simple targets that residue is small enough to absorb into the single lowering pass. For hardware whose compute model needs explicit physical state (Tenstorrent: three RISC-V threads, CB FIFOs, DST locks, SFPU LREGs), the residue becomes its own typed physical IR — see Codegen and the Physical IR.

Current Backends

Every backend except WGPU and Tenstorrent is compiled in unconditionally.

CUDA: micro-batched submission

Each CUDA device is owned by a worker thread behind a command channel. Launches accumulate in a pending window (MICRO_BATCH_WINDOW = 100); the batched-submission algorithm then distributes the whole window over per-device hardware queues (CUDA streams, default 12, config cuda.queues) with stream-wait dependencies computed between them, and submits the window at once. Events exist only inside this worker — for queue dependencies and timing (launch_timed) — never as user-facing or runtime-level objects. Buffers carry a host-side refcount and are freed behind all their in-flight work.

Configuration

Backends are configured through the process-wide config file — see Configuration. Each backend has its own section; missing sections mean that backend’s defaults.

Codegen and the Physical IR

The kernel IR’s two jobs are both value-shaped: lower movement ops into index arithmetic, and be the space the autotuner searches over — tens of thousands of variants, measured time decides. Everything decision-shaped happens before codegen. What remains for the backend is purely machine-shaped: registers, sync, placement, then bytes.

This chapter describes how codegen is organized, and the Tenstorrent physical IR (TTIR) that exists because that hardware needs more than a straight lowering.

The Two Boundary Tests

Two tests compose into the full boundary between kernel IR and backend code:

The search test — which IR holds an op

An op belongs in the kernel IR iff the autotuner should ever enumerate over it — iff there exists a variant of the program where it is placed or structured differently, and the hardware decides between them.

Movement lowering, loop structure, split/merge/pad, register tiling, local staging, tensorization, recipes: search dimensions, kernel IR. Locks, syncs, reconfigs: their placement is invariant under restructuring, nothing to search — they cost the search nothing and belong to the physical IR, placed by derivation from loop structure and checked by verify.

Consequences:

• The search space is closed. The autotuner never sees a physical IR, and no backend can smuggle a decision out of the search. If a backend-specific dimension is ever worth searching, it enters as a kernel IR config + MakeOpt, keeping the cost model measured-time-only and backend-agnostic.
• The physical IR’s no-restructuring rule is a definition, not a discipline. Restructuring is a search dimension; search lives in the kernel IR; a physical-IR op that wants restructuring failed the search test and is in the wrong IR.
• Criterion for future disputes: “would the autotuner enumerate over this?” Yes → kernel IR + MakeOpt. No → physical IR, derived placement, structural verify.

The subject test — where a pass lives

A pass belongs in an IR only if its subject exists in that IR.

If a pass reasons about values, it belongs in the kernel IR (or above). SSA machinery — rewrites, solvers, def-use analysis — is never duplicated in a physical IR. When a physical-IR pass needs a dataflow fact (e.g. “is this operand’s last use?”), it queries the kernel IR through the origin op recorded at lowering time.

The search test decides where an op lives; the subject test decides where a pass lives.

Straight Codegen Is the Default

For most backends (C, CUDA, OpenCL, Vulkan, WGPU), codegen is trivial by design: deSSA, then one linear pass over the kernel IR emitting target code. No further optimizations, no backend-specific lowering. The hard work happened in the optimization passes; the IR was deliberately shaped so backends stay thin.

Autograd

Zyx implements automatic differentiation through an explicit Tape. Unlike other frameworks, there is no separate autograd graph — it uses the same graph as computation.

In most frameworks, autograd requires a separate graph because the eager execution engine discards intermediate results. Zyx’s tape keeps all graph ops alive until realize() or drop. That is zyx keeps all ops, with each of the ops taking 32 bytes. Purely symbolic, no allocation, no computation. Even the ops themselves rarely allocate, which means you get very fast graph retracing for validity checks (measured llama retracing with some 7k tensor ops at sub 1ms per forward pass). As you might have noticed by the op count, these ops are highly granular, there is no matmul.

Tape API

extern crate zyx;
use zyx::{DType, Tape, Tensor, ZyxError};
fn main() -> Result<(), ZyxError> {
let x = Tensor::randn([2, 3], DType::F32)?;
let y = Tensor::randn([2, 3], DType::F32)?;
let tape = Tape::new([&x, &y])?;
let z = x.relu() * y.tanh();

let grads = tape.gradient(&z, vec![&x, &y]);
// grads[0] = gradient of z w.r.t. x
// grads[1] = gradient of z w.r.t. y
Ok(())
}

No “requires_grad”, no “torch.no_grad()”

There’s no requires_grad flag on tensors. The tape records the entire graph; when you call gradient(), you specify which tensors you want gradients for. There is no no_grad either. Everthing recorded on the tape is differentiable (ofc. excluding mathematically non-diff ops). The point is that autograd is just a pass that reads the graph and appends new nodes to it. The ops appended are the same variants as the ones in the forward pass, which makes optimization and fusion across forward/backward easy and it also means the graph is endlessly differentiable and higher order derivatives are supported naturally.

extern crate zyx;
use zyx::{DType, Tape, Tensor, ZyxError};
fn main() -> Result<(), ZyxError> {
let x = Tensor::randn([2, 3], DType::F32)?;
let y = Tensor::randn([2, 3], DType::F32)?;
let tape = Tape::new([&x, &y])?;
let z = y.exp();

let grads = tape.gradient(&z, vec![&x, &y]);  // grads[0] is zero — z doesn't depend on x
Ok(())
}

This is more flexible — you don’t need to decide at tensor creation time which tensors will be differentiated. Gradient sources must be part of the tape’s graph: tensors registered via Tape::new or produced by ops inside the tape.

Higher-Order Derivatives

Autograd is an append-only transform on the tape: gradient() doesn’t consume or rewrite anything — it appends backward ops to the same graph, and the gradients it produces are ordinary tensors on the tape. Higher-order derivatives are just another append: differentiate the gradient like any other tensor.

extern crate zyx;
use zyx::{DType, Tape, Tensor, ZyxError};
fn main() -> Result<(), ZyxError> {
let x = Tensor::randn([2, 3], DType::F32)?;
let tape = Tape::new([&x])?;
let z = x.relu();
let g = tape.gradient(&z, vec![&x]);
let h = tape.gradient(&g[0], vec![&x]);  // second derivative
Ok(())
}

Many properties of the autograd system fall out of this fact:

• No separate backward graph — backward ops live in the same graph as forward ops, so they get the same treatment: CSE, fusion, and kernel caching apply to backward code for free.
• Gradients are first-class tensors — they can feed further computation, be stored, or be differentiated again.
• Nothing is snapshotted or frozen — since the transform only appends, the tape needs no copy of the forward graph and no special “backward mode”.
• Memory behavior is uniform — graph ops accumulate until realize() or drop, whether they came from forward ops or backward ops.

Memory Efficiency

Inside a tape, intermediate tensors needed for backpropagation are not held in memory until realize time. The tape stores only TensorId (u32) values — not the actual data. When the tape is dropped, all tape-preserved ops are released.

Module System

The module system provides a way to group tensors (parameters) into neural network layers. It’s defined by the Module trait and powered by #[derive(Module)].

The Module Trait

pub trait Module {
    fn iter(&self) -> impl Iterator<Item = &Tensor>;
    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor>;
    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)>;
    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)>;
    fn set_params(&mut self, params: &mut HashMap<String, Tensor>);

    fn save(&self, path: impl AsRef<Path>) -> Result<(), ZyxError>;
    fn save_numpy(&self, path: impl AsRef<Path>) -> Result<(), ZyxError>;
}

#[derive(Module)]

The #[derive(Module)] macro (from zyx-derive) generates the trait implementation, collecting all tensor fields recursively. This works with nested modules:

extern crate zyx;
extern crate zyx_nn;
use zyx::Tensor;
use zyx_nn::Module;
#[derive(Module)]
struct Linear {
    weight: Tensor,
    bias: Tensor,
}

#[derive(Module)]
struct MLP {
    layer1: Linear,
    layer2: Linear,
    layer3: Linear,
}
fn main() {}

Using Modules

extern crate zyx;
extern crate zyx_nn;
extern crate zyx_optim;
use zyx::{DType, Module, Tape, Tensor, ZyxError};
use zyx_nn::Linear;
use zyx_nn::Module;
use zyx_optim::SGD;
#[derive(Module)]
struct SimpleNet {
    linear1: Linear,
    linear2: Linear,
}
impl SimpleNet {
    fn new() -> Result<Self, ZyxError> {
        Ok(Self {
            linear1: Linear::new(4, 8, true, DType::F32)?,
            linear2: Linear::new(8, 2, true, DType::F32)?,
        })
    }
    fn forward(&self, x: &Tensor) -> Tensor {
        let x = self.linear1.forward(x).unwrap().relu();
        self.linear2.forward(&x).unwrap()
    }
}

fn train_step(model: &mut SimpleNet, optim: &mut SGD, x: &Tensor, target: &Tensor) -> Result<f32, ZyxError> {
    let tape = Tape::new(&*model)?;
    let output = model.forward(x);
    let loss = output.mse_loss(target)?;
    let grads = tape.gradient(&loss, &*model);
    optim.update(&mut *model, grads);
    tape.realize(model.iter().chain([&loss]))?;
    // tape drop cleans up graph state; next iteration reuses cached plan
    Ok(loss.item())
}
fn main() -> Result<(), ZyxError> {
    let mut model = SimpleNet::new()?;
    let mut optim = SGD::default();
    let x = Tensor::randn([2, 4], DType::F32)?;
    let target = Tensor::randn([2, 2], DType::F32)?;
    let _ = train_step(&mut model, &mut optim, &x, &target)?;
    Ok(())
}

The tape.gradient(&loss, &model) call passes the model itself as the sources. The autograd system iterates over model.iter() to get all parameters.

Serialization

Modules can save and load parameters in safetensors format:

extern crate zyx;
extern crate zyx_nn;
use zyx::{DType, Module, Tensor, ZyxError};
use zyx_nn::Linear;

#[derive(zyx_nn::Module)]
struct SimpleNet {
    linear1: Linear,
}

fn main() -> Result<(), ZyxError> {
let mut model = SimpleNet { linear1: Linear::new(4, 2, true, DType::F32)? };
model.save("model.safetensors")?;
let mut params = Tensor::load_safetensors("model.safetensors")?;
model.set_params(&mut params);
Ok(())
}

Runtime and Plans

The runtime (runtime.rs) is the process-wide state of zyx: graphs, tensor states, kernels, compiled programs, measured timings, interned shape expressions, and the plan cache. Memory pools belong to the backends, not to the runtime — they outlive it.

The Runtime

pub struct Runtime {
    pub graphs: Slab<GraphId, Graph>,
    pub tensors: Slab<TensorId, TensorData>,
    pub kernels: Slab<KernelId, KernelData>,
    kernel_map: Map<Kernel, KernelId>,          // content-dedup of kernels
    programs: Map<KernelId, DeviceProgramId>,   // compiled programs
    timings: Map<ProgramId, u64>,               // measured nanos
    pub(crate) expr_hash: Map<Expr, ExprId>,    // hashcons of shape expressions
    pub exprs: Slab<ExprId, Expr>,              // append-only interned exprs
    pub rng: Rng,
    pub implicit_casts: bool,
    pub training: bool,
    pub plan_cache: Map<u64, Plan>,
    #[cfg(feature = "viz")]
    pub viz: Viz,
}

The runtime is stored in a global mutex:

static RT: mutex::Mutex<Runtime> = mutex::Mutex::new(Runtime::new());

Every tensor operation locks it. The lock is coarse: it covers graph mutation, and also the whole realize path — Tape::realize holds it while compiling the plan, replaying launches, and converting outputs to eager tensors. It is not a fine-grained per-tensor lock, and nothing re-enters it (helpers that need the runtime take &mut Runtime rather than locking again). Backend pools and their device worker threads are separate, process-wide globals.

Shape expressions are interned: ExprIds are stable forever, structural equality is one ExprId compare, and symbolic dims never live in the tensors slab. Tensor handles are u32 indices into tensors.

Tensor States

A TensorData is one of:

item() forces an eager value: the producer kernel is materialized and launched. shape() does not force; symbolic shapes are resolved without allocating.

Memory Pools and Placements

Pools are process-wide and never deinitialized:

pub enum Pool {
    Host,          // RAM, shared by C and Host
    Disk,          // file-backed chunks (paths, not bytes)
    Cuda(u16),
    OpenCL(u16),
    Vulkan(u16),
    TT(u16),       // feature = "tenstorrent"
    WGPU(u16),     // feature = "wgpu"
}

Each backend owns its globals — one mutex per pool ordinal — and lazily initializes them on first use. A value is a Placement:

pub struct Placement { pub shards: Vec<Shard> }

pub struct Shard {
    pub pool: Pool,
    pub chunk: ChunkId,
    pub offset: usize,
    pub len: usize,
}

One shard is the common case; several shards describe a value split across devices. Placement’s Drop returns each chunk to its pool’s free list; Pool::release frees nothing by itself.

Devices and pools are two views of the same hardware: Dev::pool() derives the pool from the device, never the reverse.

Device Selection

Runtime::pick_device(bytes) is the generic placement rule:

1. start from Dev::all() — lazy init of every compiled-in backend, cached in a OnceLock; backends that are configured out or whose hardware/driver is missing contribute nothing;
2. drop AOT-only devices (Dev::aot_only()) for generic kernels;
3. keep devices whose pool has at least bytes free;
4. sort by free_compute() descending and take the first.

If no device has enough free memory the call is an allocation error. A user can pin a device with ZYX_DEV=cpu|cuda|vulkan|... (parsed by Dev::from_str) or by calling Tensor::to(Dev) explicitly; Dev::Auto resolves through the rule above.

Plans

Every launch is a Plan — a scheduled CmdQueue:

pub enum Cmd {
    Launch { program: ProgramId, args: Vec<OpId>, outputs: Vec<(OpId, DType, Vec<PlanDim>)> },
    Copy   { src: OpId, dst: OpId, dst_pool: Pool, dst_dtype: DType, dst_dims: Vec<PlanDim> },
    Alias  { class: OpId, to: OpId },
}

Slots are queue-local OpIds. PlanDim is a symbolic output size built from constants and leaf classes; it is evaluated at replay against class_vars.

CmdQueue::schedule(outputs):

• slices the queue into maximal same-device launch runs; each run is scheduled by its own device;
• executes every Copy on the destination device (device runs never copy);
• joins Alias bindings to the surrounding run as zero-cost map entries;
• pins each run’s outputs plus the slots read after it, and frees everything else at last use.

Plan::replay(boundary, class_vars) binds the boundary table to Placements, allocates every unbound output from its PlanDims, resolves aliases through the boundary, and returns the placements of the plan’s outputs. No allocation or liveness decision was made at schedule time beyond this structure.

The plan cache

Runtime::plan_cache is keyed by the graph’s content hash (structure + outputs) folded with the pool of every leaf class at call time. Two realizations of the same graph shape share a plan only when their leaf pools match; otherwise the plan recompiles. This is what makes a tape loop cheap after the first iteration: the kernels are already compiled and the plan is already scheduled, so later iterations only replay it with new leaf values and (possibly) new variable dims.

Eager execution

An eager kernel follows the same route with a one-kernel queue. On materialization the runtime:

1. verifies argument arity;
2. linearizes, runs CSE/DCE/instruction scheduling;
3. merges group indices down to the device’s max_global_work_dims and renumbers them;
4. verifys;
5. runs delete_useless_loops and three rounds of default_epilogue;
6. runs BeamSearch (or NoSearch under ZYX_DEBUG=256);
7. stores the winning program and its measured time.

The resulting plan is executed immediately. Graph kernels go through the same compile path via autotune_jit_kernels, but their programs are recorded as Op::Program nodes and extraction decides whether they run.

Async Execution

Synchronization is per backend and lives below Plan. A plan replay issues launches and copies in program order; each backend decides how far it can run ahead. CUDA, for example, owns each device with a worker thread behind a command channel: launches accumulate in a micro-batch window, are distributed over several hardware queues with computed stream dependencies, and are submitted together. There are no runtime-level event objects — ordering is expressed by the plan (After, copy direction, argument dependencies) and enforced by the backend worker.

Lazy Device I/O

A tensor can describe data that still lives on disk. Runtime::new_disk_tensor creates a Leaf whose shard is a Pool::Disk chunk holding a path and an offset instead of bytes. When a plan needs the value on a compute device, the copy command stages it there; the disk pool keeps the file offset information, so nothing is read until the value participates in a launch.

The eager loader is Module::load_safetensors(path), which reads a safetensors file into a HashMap<String, Tensor>; disk-backed leaves are the mechanism underneath for deferring that read.

Lifecycle

Runtime::deinitialize runs at process exit. With the time feature it prints the accumulated timing table first, then drops the tensor and kernel slabs. Backend pools are static and are not deinitialized — the process owns the hardware for its lifetime. Runtime::manual_seed reseeds the runtime RNG.

Configuration

Zyx is configured through a JSON file and environment variables. The file is read once per process, on the first backend access; changing it requires a restart. Missing sections and missing keys fall back to defaults, and a missing or unparsable file simply means “all defaults”.

Config File

The runtime reads $XDG_CONFIG_HOME/zyx/config.json when XDG_CONFIG_HOME is an absolute path, otherwise ~/.config/zyx/config.json. ZYX_DEBUG=1 prints whether the file was found and parsed.

A complete example (every section is optional):

{
    "autotune": {
        "n_launches": 1,
        "n_seeds": 200,
        "n_added_per_step": 200,
        "n_removed_per_step": 0,
        "n_total_opts": 1000
    },
    "c": { "enabled": true },
    "cuda": { "device_ids": [0], "cudnn": true, "queues": 8 },
    "opencl": { "platform_ids": [0], "queues": 8 },
    "vulkan": { "device_ids": [0] },
    "tenstorrent": { "device_ids": [0] }
}

wgpu ("wgpu": { "enabled": true }) exists only when the crate is built with the wgpu feature; the tenstorrent section only has an effect when the crate is built with the tenstorrent feature.

Backend Sections

device_ids / platform_ids: null/absent means “every device the backend finds”; a list selects exactly those ordinals; an empty list disables the backend. Ordinals are the backend’s own (CUDA driver order, Vulkan device index, OpenCL platform index).

cuda.cudnn enables the cuDNN AOT matmul path (the struct default is false). cuda.queues and opencl.queues set the number of hardware command queues/streams per device used by the submission workers (default 8).

Autotune Section

See Optimization Passes for how the search uses them.

Environment Variables

ZYX_DEV — device pinning

ZYX_DEV overrides the config file and leaves exactly one device available; every other backend is configured out (and a device with an ordinal keeps only that ordinal). It accepts the same syntax as Dev::from_str:

ZYX_DEV=cpu          # C backend
ZYX_DEV=cuda:0        # CUDA driver ordinal 0
ZYX_DEV=tt            # Tenstorrent, device 0 (feature "tenstorrent")
ZYX_DEV=vulkan        # Vulkan device 0
ZYX_DEV=opencl:1      # OpenCL platform 1

Names are case-insensitive: auto, c/cpu, host, cuda, tt/tenstorrent, vulkan/vk, opencl/cl, wgpu. An index on an indexless device (c:1) is an error. A bad value panics — silently falling back would run on the wrong hardware — and ZYX_DEV=auto is a no-op (no narrowing).

ZYX_DEBUG — trace flags

A bitmask; add the values you want:

ZYX_DEBUG=8 cargo run     # print the winning kernel IR
ZYX_DEBUG=24 cargo run    # IR (8) + generated code (16)
ZYX_DEBUG=256 cargo run   # NoSearch: compile a cleaned seed once

Other variables

API Configuration

Tensor::manual_seed(42);        // reseed the runtime RNG
Tensor::set_training(true);     // enable training mode (autograd graph is kept)
Tensor::set_training(false);    // inference
let t = Tensor::training();     // current mode

Device pinning can also be done per tensor with Tensor::to(Dev); see Backend System.