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 no more fusion is beneficial (based on a heuristic), the kernel executes immediately. 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 (egraph-based fusion comparison, device allocation search, 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 creates a graph op. 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, egraph explores
                     fusion variants + device allocations)
                                  │
                                  ▼
                    Kernel IR ──► Opt ──► Codegen ──► Execute

Inside a tape, graph ops accumulate lazily. At realize time, the kernelizer processes the full graph while egraph exploration tries different fusion schemes and device allocations, selecting the fastest. 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”), egraph exploration of fusion variants, device allocation search, 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.

Architecture Overview

The zyx pipeline transforms high-level tensor operations into device-specific machine code.

Tensor API ──► Eager-ish: append to kernel ──► compile + execute
            │
            └── Tape: Graph (lazy) ──► Kernelizer ──► Kernel IR ──► Opt Passes ──► Backend Codegen
                                                                                    │
                                       Autotune (clone + evaluate)  └── deSSA + linear pass

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(())
}

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 — no graph, no separate realize step. Inside a Tape, ops build graph ops lazily. Either way, each operation returns a lightweight Tensor handle.

2. The Graph

The graph opset was taken from tinygrad, with changes to make it even smaller. This is the minimal set of operations that can express ALL linear algebra operations and ALL PyTorch ops — by stacking these ops:

Tensor handles are u32 (4 bytes) into the slab. Still small enough that 10,000 handles cost ~40 kB.

The graph is stored in a Slab — a dense array with free-list tracking. TensorId is a u32 index into this slab, making tensor handles 4 bytes.

3. The Kernelizer

The kernelizer fuses compatible graph ops into kernels. Outside a tape it runs incrementally as ops are added; inside a tape it runs when Tape::realize() is called (dropping only cleans up graph state). The kernelizer uses heuristics to decide where kernel boundaries go — it’s not a simple rule. A reduce op used by multiple downstream ops does not necessarily force a split. If two downstream ops are both expand ops, that may force fusion. Element-wise chains will almost always fuse into one kernel.

View operations (reshape, expand, permute, pad) are unfolded into index arithmetic in the kernel, becoming “free” — they don’t create separate operations.

4. The Kernel IR

After unfolding, the DAG is converted into a linear structure — a doubly-linked list of ops stored in an arena (the Slab<OpId, OpNode>). Each OpNode is 32 bytes stored inline in the arena — no Box, no vtables, no indirection.

OpId is a u32 index into the slab — random access is O(1). The IR is SSA, except for loops and Define ops (which can be mutable).

The design goal of the small opset: optimizations are easy to write. If you understand the IR, you can add a new pass in an afternoon.

5. Optimization Passes

Optimization passes work on the linear IR. Kernels are cloned and each variant is evaluated separately — no egraphs. The cost function evaluates thousands of variants per second.

Optimization passes emit IR that is deliberately simple to lower to backend instructions. A backend just does deSSA + a single linear pass over the IR. No complex backend-specific lowering.

The autotune system explores the search space by:

1. Starting with the initial kernel
2. Applying one optimization variant
3. Hashing the kernel to detect duplicates
4. Evaluating with the cost function (or launching and timing)
5. Repeating by combining optimization sequences
6. Selecting the best variant

There is no fixed number of optimization passes zyx aims to have. The design goal was a small opset in the graph and IR so that passes are easy to write. More passes will be added over time.

6. Backend Codegen

Each backend converts the stabilized kernel IR into target code. Since the IR is designed for this, codegen is a straight line: deSSA, then one pass over the ops emitting instructions. No further optimizations, no complex lowering. Hardware whose compute model can’t be expressed this way (Tenstorrent) gets an explicit physical IR — see Codegen and the Physical IR.

Backends are dispatch via enums (no dyn Backend — that would require downcasting, which is ugly in Rust):

pub enum Device {
    C(CDevice),
    CUDA(CUDADevice),
    OpenCL(OpenCLDevice),
    Vulkan(VulkanDevice),
    WGPU(WGPUDevice),
    HIP(HIPDevice),
    // Tenstorrent, etc.
}

All backends are compiled into the library and selected at runtime.

7. Runtime and Scheduler (Current)

The scheduler picks a device based on free memory and compute capacity. Cross-device data transfers are scheduled as kernels; asynchronous execution is managed per backend (e.g. CUDA’s micro-batched submission over hardware queues — see Backend System).

Debugging the Pipeline

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, so any dim expression can be fully evaluated to a concrete constant at any time: walk the tree, fold Const leaves via Constant::unary / Constant::binary, and read variable slots at Param { Variable } leaves. Consumers must evaluate dims this way instead of fabricating placeholder values (0, -1, …) where evaluation can produce the real value.
• Inline ops — all ops live in the arena as flat 32-byte entries. No Box, no vtables, no indirection. Passes allocate their own working data (hash maps, vecs) as needed.
• Linear IR — linked list of fixed-size ops. Optimizations traverse front-to-back or back-to-back.
• Backend codegen is trivial — the hard work is in the IR-level optimization passes.
• Tape-scoped lazy graph — ops inside a tape build graph ops instead of executing eagerly. Enables egraph fusion, device allocation search, 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 graph 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, all metadata lives in the graph, not the tensor 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 (lazy loading).

The Immutability Rule

Tensors are immutable — there is no in-place mutation:

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(())
}

This makes autograd simpler (no mutation to track) and eliminates backpropagation errors from in-place modifications.

The Graph

The graph is an e-graph (equivalence graph) used in tape mode for tensor operation rewrites, autograd, fusion and device placement and movement decisions. 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 {
    hashcons: Map<Op, OpId>,
    ops: Slab<OpId, OpNode>,
    jit_kernels: Slab<JitKernelId, JitKernelData>,
    leaf_classes: Vec<OpId>,
    leaf_map: Map<OpId, TensorId>,
    ref_count: u64,
    dead: bool,
    max_cons_id: u32,
}

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. Each op belongs to an equivalence class identified by its class_of field (the root OpId of the class). Equivalent forms of the same computation (e.g. different layouts of a matmul) live in the same class.

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, cons_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]> },
    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 },
    After { x: OpId, dep: OpId },
    ToDevice { x: OpId, device: Dev, time: u64 },
    Contiguous { x: OpId },
    Program { inputs: OpId, outputs: OpId, info: Box<(ProgramId, u64)> },
    Kernel(Box<CustomKernel>),
    ...
}

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; Bitcast reinterprets bits without value conversion; After orders side-effecting ops; Contiguous materializes a layout; Custom wraps opaque custom kernels.

The reason to use OpId instead of TensorId is to get cannonicalization 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. The graph supports rewrites that produce equivalent forms of a computation:

• CSE via hashconsing
• Algebraic rewrites like transpose fusion
• Layout rewrites: matmul can be realized as transposed or un-transposed
• Shape rewrites: reshape and padding can be fused or split

There is no cost model: each fusion variant is individually autotuned, and Graph::extract picks 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 10,000 handles cost ~40 kB. Op is 32 bytes, so 10k ops costs 320kB. When the tape is dropped, the graph is deleted, so there is no more memory being eaten by anything. So there is no gc and no refcounting inside graph, it’s append only for the most part and deleted as a whole (arena style).

The Kernelizer

The kernelizer (kernelize.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, the kernelizer runs incrementally as each op is added — fusing compatible nodes and executing when a fusion boundary is hit. Inside a tape, the kernelizer runs once during Tape::realize() (dropping only cleans up graph state):

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

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(). 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?

Stores also trigger automatically when a graph node is requested as the final output (to_eval), creating a natural boundary there.

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

struct Kernelizer<'a> {
    must_keep_nodes: Set<TensorId>,
    pending_stores: Set<TensorId>,
    realized_nodes: Set<TensorId>,
    kernels: Slab<KMKernelId, Kernel>,
    visited: Map<TensorId, (KMKernelId, OpId)>,
    rcs: Map<TensorId, u32>,
    graph: &'a Graph,
    // ...
}

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

pub struct KernelCache {
    pub cache: Map<KernelId, Kernel>,
}

After a kernel is built and optimized, its hash is computed. If the same kernel was compiled before, the cached compiled program is reused. The cache persists across realize() calls, so repeated graph patterns (like training loop iterations) hit the cache on the second iteration.

The Kernel IR

The kernel IR is the intermediate representation used for all computation kernels. It is a doubly-linked list of 32-byte OpNodes stored in an arena (the Slab allocator).

Data Structure

pub struct Kernel {
    pub ops: Slab<OpId, OpLinked>,
    pub head: OpId,
    pub tail: OpId,
    // ...
}

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

The Slab<OpId, OpNode> is a Vec<OpNode> with a free-list. OpId is a u32 index — random access is O(1).

Linearization

Before linearization, view ops (Reshape, Expand, Permute, Flip, Pad, Narrow) are represented directly as Op variants. They are unfolded into index arithmetic during the linearization pass — after linearization, all ops are fixed-size inline entries in the arena — no Box, no vtables, no per-op indirection.

The IR is in SSA form, except for Loop, Param, and Storage ops (which can carry mutable state).

Op Variants

Parameters

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

Arithmetic

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::Mad { x: OpId, y: OpId, z: OpId }
Op::Stack { ops: Box<[OpId]> }

Memory

Op::Storage { dtype, scope: MemScope, len: Dim }  // kernel-internal memory
Op::Load { src, index, layout }
Op::Store { dst, src, index, layout }
Op::Const(Constant)

Op::Param and Op::Storage are the SSA escape hatches: the mutable stuff values are read from and written to across the linear order.

Control Flow

Op::Loop { len: OpId }
Op::EndLoop
Op::If { condition: OpId }
Op::EndIf
Op::Barrier

Indexing

Op::Range { axis, kind: RangeKind }  // Group / Local / Scalar
Op::Index { vec: OpId, idx }         // select a value from a Stack

Hardware Accelerators / Tiles

Op::Wmma { dims, layout, dtype, a, b, c }

Backend-Specific

Op::Asm { .. }  // inline assembly for backends with JIT asm (e.g. Tenstorrent)
Op::GPU { .. }
Op::TT { .. }
Op::PTX { .. }

View 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 }
Op::Reduce { x: OpId, rop: BOp, reduce_axis: OpId }

Memory Layouts and Scopes

pub enum MemLayout {
    Scalar,
    Vector(u8),
    Tile { x, y, stride },
}

pub enum MemScope {
    Global,
    Local,
    Register,
    Variable,
    Circular,
}

Backend Codegen

Because the IR is designed for it, backend codegen is trivial:

1. deSSA — resolve SSA references to physical registers/memory
2. Linear pass — walk the op linked list once, emitting instructions

No further optimizations, no complex lowering. Backends whose compute model needs explicit physical state (Tenstorrent) get an additional physical IR below the kernel IR — see Codegen and the Physical IR.

Debugging

Set ZYX_DEBUG=8 to print the kernel IR:

r18: i32 = def global, len=4
r44: u32 = gidx0    // 0..=0
r19: i32 = r18[r1]  // 0..=3 load

Optimization Passes

Optimization passes operate on the linear kernel IR. They are divided into always-on passes (run before every compilation) and autotuned passes (searched at runtime for the best variant).

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 will be added over time as performance opportunities are identified.

Always-On Optimizations

These run in a fixed pipeline (default_epilogue) on every kernel state during autotuning and before compilation:

pub fn default_epilogue(&mut self) {
    self.unroll_len1_loops();
    self.constant_folding();
    self.move_constants_to_beginning();
    self.loop_invariant_code_motion();
    self.fold_accs();
    self.delete_zero_len_indices();
    self.delete_zero_len_loops();
    self.unfold_pows();
    self.algebraic_simplifications();
    self.simplify_accumulating_loop();
    self.swap_commutative();
    self.common_subexpression_elimination();
    self.instruction_schedule();
    self.dead_code_elimination();          // must stay last
    // + exp ↔ exp2 conversion (per device capability)
    // + tenstorrent tile lowering when targeting TT
}

Constant Folding

Evaluates expressions where all inputs are constants. For example, 2 + 3 becomes 5 and sin(0.0) becomes 0.0.

Motion of Constants to Beginning

Moves all constant definitions to the start of the kernel. This creates better fusion opportunities and simplifies register allocation.

Loop-Invariant Code Motion (LICM)

Hoists operations that produce the same value on every iteration to before the loop.

Algebraic Simplifications / Unfold Pows / Swap Commutative

Rewrites like x*1 → x, x+0 → x, exponentiation unfolding, and normalizing commutative operand order so CSE’s hashing sees identical forms.

Accumulator Folding / Simplify Accumulating Loop

Simplifies accumulator update patterns and loop-carried reductions for more efficient code generation.

Instruction Scheduling

Orders ops within basic blocks to improve ILP before codegen.

Dead Code Elimination

Removes ops whose results are never used. This is always the final pass — backends fail on unreferenced ops.

Autotuned Passes

The autotune system (BeamSearch) clones the kernel, applies optimization variants, and evaluates each separately. No egraphs — just clone, transform, hash, evaluate. The cost function can evaluate thousands of variants per second.

pub const fn default_optimizations() -> [MakeOpt; 8] {
    [
        Kernel::opt_split_global_to_local,
        Kernel::opt_reassociate_commutative,
        Kernel::opt_coarsen,
        Kernel::opt_register_blocking,
        Kernel::opt_local_reduce,
        Kernel::opt_split_loop,
        Kernel::opt_merge_nested_loops,
        Kernel::opt_fuse_mad,
    ]
}

Split Global to Local

Adjusts block and thread dimensions for better memory access patterns. For example, a kernel with block_dim = 1024, thread_dim = 1 becomes block_dim = 32, thread_dim = 32. This enables coalesced memory access on GPU.

Reassociation

Reorders commutative operations to create more fusion opportunities.

Coarsen

Merges parallel work per thread so each thread handles multiple elements (the inverse of splitting).

Register Blocking

Unrolls tree reductions and coarsens global threads so each thread processes multiple elements, increasing computational intensity and register reuse.

Local Reduce

Reduces within a thread’s local data before hitting global or shared memory.

Loop Splitting

Splits large loops into chunks for better register pressure and instruction-level parallelism.

Merge Nested Loops

Collapses nested loop nests into a single loop where semantics allow.

Fuse Multiply-Add

Fuses mul + add chains into Op::Mad.

Performance Budget

Every pass must stay fast — tens of microseconds is the ceiling, single-digit microseconds the norm. Passes run once per autotune variant and per compiled kernel, so slow passes compound into seconds of autotuning.

How Autotuning Searches

1. Start with the initial kernel and run the epilogue
2. Apply ONE optimization variant and run the epilogue
3. Hash the kernel — skip if already visited
4. Launch and measure, or evaluate with the cost function
5. Repeat by combining with existing optimization sequences
6. Select the best variant; only the top variants are compiled and launched

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. The verify pass (verify.rs) checks internal IR consistency in debug mode.

Backend System

Zyx supports multiple hardware backends. Backends are enum-dispatched, compiled into the library (feature gates for WGPU and TT), 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,
    /// CBLAS backend for AOT matmuls (runs on Pool::Host).
    Cblas,
    /// 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 CBLAS 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, CBLAS: single devices, 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

All backends except WGPU and Tenstorrent are compiled in by default; those two require --features wgpu / --features tenstorrent.

CBLAS

The CBLAS device runs only AOT (precompiled) matmul kernels and cannot compile generic zyx kernels. Dev::aot_only() marks it, and generic kernel autotuning skips it — it competes only where a precompiled BLAS call is an option.

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/coarsen/register blocking, MAD fusion: 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, HIP, 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.

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

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 realize(&self) -> Result<(), ZyxError>;
    fn save(&self, path: impl AsRef<Path>) -> Result<(), ZyxError>;
    fn set_params(&mut self, params: &mut HashMap<String, Tensor>);
}

#[derive(Module)]

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

#[derive(Module)]
struct Linear {
    weight: Tensor,
    bias: Tensor,
}

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

Using Modules

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

fn train_step(model: &mut SimpleNet, optim: &mut SGD, x: &Tensor, target: &Tensor) -> f32 {
    let tape = Tape::new(&*model)?;
    let output = model.forward(x);
    let loss = output.mse_loss(target)?;
    let grads = tape.gradient(&loss, &*model);
    optim.update(model, grads);
    // tape drop cleans up graph state; next iteration reuses cached plan
    loss.item()
}

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:

model.save("model.safetensors")?;
let params = Tensor::load_safetensors("model.safetensors")?;
model.set_params(&mut params);

Runtime and Scheduler

The runtime (runtime.rs) is the global state of zyx — the graph, devices, buffers, and kernel cache all live here.

The Runtime

pub(crate) struct Runtime {
    pub graph: Graph,
    pub devices: Slab<DeviceId, Device>,
    pub pools: Slab<PoolId, MemoryPool>,
    pub buffer_map: Map<TensorId, BufferId>,
    pub events: Map<BTreeSet<BufferId>, Event>,
    pub kernel_cache: KernelCache,
    pub rng: Rng,
    pub autotune_config: AutotuneConfig,
    pub debug: DebugMask,
    pub training: bool,
}

The runtime is stored in a global Mutex<Runtime>:

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

Every tensor operation locks this mutex, appends a graph node (microseconds), and releases. The lock is never held during computation — only during graph manipulation.

Memory Pools

Each backend provides a MemoryPool for allocating device buffers:

pub enum MemoryPool {
    Host(HostMemoryPool),
    Disk(DiskMemoryPool),
    C(CMemoryPool),
    CUDA(CUDAMemoryPool),
    // ...
}

The Scheduler (Current)

The current scheduler (schedule.rs) selects a device for kernel execution by calculating required memory, sorting devices by free compute capacity, and picking the first with enough free memory. It handles cross-device transfers via events.

New Scheduler (In Development)

A new scheduling approach is under development in search.rs and search2.rs. It will use an e-graph-like budget-guided exhaustive fusion enumeration, including costs for memory movement operations — replacing the current simple heuristics with exploration of all fusion configurations within a cost budget.

Async Execution

Events track kernel completion:

events: Map<BTreeSet<BufferId>, Event>

Lazy Device I/O

Tensors can reference data on disk without loading it:

let t = Tensor::from_safetensors("model.safetensors", "layer.weight")?;

The disk pool keeps file offset information. Data is loaded lazily when the tensor needs to be realized on a compute device.

Configuration

Zyx is configured through a JSON file, environment variables, and API calls.

Config File

The runtime reads $XDG_CONFIG_HOME/zyx/config.json (or ~/.config/zyx/config.json by default):

{
    "c": { "enabled": true },
    "cuda": { "device_ids": [0] },
    "opencl": { "platform_ids": [0] },
    "autotune": {
        "enabled": true,
        "max_configs": 100,
        "timeout_ms": 5000
    }
}

Backend Control

Environment Variables

Debug Flags

Set ZYX_DEBUG as a bitmask:

ZYX_DEBUG=8 cargo run     # print kernel IR
ZYX_DEBUG=16 cargo run    # print generated code

Colorless Output

AGENT=1 ZYX_DEBUG=8 cargo test

API Configuration

Tensor::manual_seed(42);            // deterministic RNG
Tensor::set_training(true);         // enable training mode
Tensor::set_training(false);        // disable training mode