Introduction

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

Unlike traditional ML frameworks that separate the eager execution graph from the autograd graph, zyx uses a single unified graph — but only inside a Tape scope. Outside a tape, there is no graph: ops go directly to the kernelizer. Inside a Tape, the graph is shared between computation and autograd, eliminating duplication and keeping the implementation lean — tensors are only 4 bytes and the graph uses ~20 node types.

Who This is For

This book is for developers who want to understand how zyx works under the hood: the architecture decisions, the optimization passes, the backend system, and how pieces fit together.

Two Execution Modes

Zyx operates in two modes:

• Eager-ish (default): Tensor operations fuse into kernels as you write them. When no more fusion is possible, 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. Kernel caching (compiled program reuse) is shared across both modes.

The Architecture at a Glance

Every tensor operation creates a graph node. What happens next depends on the mode:

Eager-ish (default) — direct fusion, no graph:

Tensor op ──► append to kernel ──► compile + execute (if fusion not possible)

Each op is appended directly to the kernel that produced its inputs. When a new op can’t fuse (different device, incompatible data flow), the pending kernel compiles and executes. No graph, no separate realize step.

Tape — lazy graph, egraph exploration:

Tensor op ──► graph node (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 nodes 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 nodes — the tape prevents their deletion so the backward pass can traverse them.

Why This Design

Most deep learning libraries use two separate graphs:

1. A compute graph for eager execution
2. An autograd graph for backpropagation

Zyx uses one graph for both. This means:

• The autograd system doesn’t need its own graph infrastructure — it reuses the same nodes
• Kernel fusion works across operation boundaries without special handling
• The implementation is debuggable (one graph to inspect, not two)
• Memory overhead is minimal: tensor handles are u32 (4 bytes)

The trade-off in tape mode is that evaluation is lazy — you must call realize() to trigger computation (dropping the tape only cleans up graph state — no computation is performed). But this laziness enables optimizations that eager execution cannot: 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 appropriate for one-off ops. Kernel caching (compiled program reuse) is shared across both modes.

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 nodes 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 nodes:

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 nodes 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 node used by multiple downstream nodes does not necessarily force a split. If two downstream nodes 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),
    Dummy(DummyDevice),
    // 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 nodes. 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 nodes 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 node, 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 nodes 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 and optimization. Inside a Tape, every operation builds a node 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<Node, NodeId>,
    nodes: Slab<NodeId, NodeData>,
    classes: Slab<ClassId, EClass>,
    ekernels: Slab<EKernelId, EKernelData>,
    kernel_map: Map<NodeId, EKernelId>,
    leaf_map: Map<ClassId, TensorId>,
    rc: u32,
    max_leaf_id: u32,
}

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

Each NodeId maps to a NodeData entry in the nodes slab, and each node belongs to an equivalence ClassId in the classes slab. Equivalent forms of the same computation (e.g. different layouts of a matmul) live in the same class.

Node Types

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

enum Node {
    Const(Constant),
    Leaf { cons_id: u32, dtype: DType, shape: ClassId },
    Expand { x: ClassId, shape: ClassId },
    Permute { x: ClassId, axes: Box<[UAxis]> },
    Reshape { x: ClassId, shape: ClassId },
    Pad { x: ClassId, axis: UAxis, lp: ClassId, len: ClassId },
    Flip { x: ClassId, axes: Box<[UAxis]> },
    Narrow { x: ClassId, axis: UAxis, start: ClassId, len: ClassId },
    Stack { ops: Box<[ClassId]> },
    Reduce { x: ClassId, rop: BOp, axes: Box<[UAxis]> },
    Cast { x: ClassId, dtype: DType },
    Bitcast { x: ClassId, dtype: DType },
    Unary { x: ClassId, uop: UOp },
    Binary { x: ClassId, y: ClassId, bop: BOp },
    Assign { dst: ClassId, src: ClassId },
    After { x: ClassId, dep: ClassId },
    ToDevice { x: ClassId, device: Dev, time: u64 },
    Contiguous { x: ClassId },
    Kernel { inputs: Box<[ClassId]>, outputs: Box<[ClassId]>, program_id: ProgramId, time: u64 },
    Custom { inputs: Box<[ClassId]>, outputs: Box<[(ClassId, ClassId, DType)]>, /* ... */ },
}

All inputs reference ClassId rather than TensorId — nodes operate on equivalence classes, not specific tensors. View nodes (Expand, Permute, Reshape, Pad, Flip, Narrow) are per-axis and stackable; Stack builds vectors; Bitcast reinterprets bits without value conversion; Assign models in-place updates; After orders side-effecting nodes; Contiguous materializes a layout; Custom wraps opaque custom kernels.

Lifecycle with Tape

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

Inside a tape, nodes 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 nodes that the tape references are preserved for autograd; unreferenced nodes are released.

Graph Size

The graph is designed to stay small. Tensor handles are u32 (4 bytes), so 10,000 handles cost ~40 kB. When the tape is dropped, the graph shrinks back to baseline.

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 nodes bottom-up in topological order. Each graph node type adds ops to the kernel its input lives in — fusion is the default, splitting happens only when needed.

Per-Node 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 nodes have been converted to kernel ops. When a graph node 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 node 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, OpNode>,
    pub head: OpId,
    pub tail: OpId,
    // ...
}

pub struct OpNode {
    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).

Unfolding

Before any optimization passes run, the kernel IR is unfolded — LoadView, StoreView, and Move ops are converted to direct index arithmetic (Load, Store with computed indices). After unfolding, 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, If, and Define 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, IDX_T)
// 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 }
Op::ReduceTile { x, .. }
Op::MatmulTile { a, b, .. }
Op::TransposeTile { x, .. }
Op::BroadcastTile { x, .. }

Backend-Specific

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

View (before unfolding)

Op::Move { x: OpId, mop: Box<MoveOp> }
Op::Reduce { x: OpId, rop, n_axes }

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 aside), 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"
    /// Testing dummy device (config-gated).
    Dummy,
}

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"
    /// Testing dummy pool (config-gated).
    Dummy,
}

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:

impl Dev {
    pub fn all() -> Vec<Dev> {
        // C, CBLAS, Dummy: 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(self, program_id: DeviceProgramId, args: &[LaunchArg]) -> Result<(), BackendError>;
pub fn launch_timed(self, program_id: DeviceProgramId, args: &[LaunchArg]) -> Result<u64, BackendError>;

plus pool operations (alloc, free, retain/release, pool_to_host, pool_to_pool) 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.

TTIR — the Tenstorrent Physical IR

The Tenstorrent (Wormhole) compute model — three RISC-V threads (reader, math, writer), circular buffers with FIFO semantics, DST tile slots with lock/commit protocol, SFPU LREGs, unpacker/packer format state — cannot be expressed as a straight per-op lowering. The legacy codegen fused lowering, scheduling, and text emission into one walk with emergent state machines (CB accounting, init placement, format trackers). TTIR makes that state explicit: a typed, ordered, physical instruction stream that passes rewrite and a verifier checks, so invalid programs never reach the JIT.

Representation

One program = one Vec<TTOp> covering all three RISC-V kernel sources. Section boundaries are ops in the stream:

[reader ops] EndReader [compute ops] EndCompute [writer ops] EndWriter

• The barriers that today exist only as scan positions become first-class ops; verify walks one stream and checks cross-section facts (CB push/pop totals settled at each boundary).
• The renderer splits at the markers — one stream in, three kernel files out.
• Ids are typed and distinct — VId (virtual, kernel OpIds at lowering), LRegId (SFPU LREGs), DstId (DST slots), VarId (scalar C registers), CBId (circular buffers) — so misuse is a compile error in the IR’s own definitions.
• Not SSA: value ops carry their target (z: OpId); SSA-ness holds as an asserted invariant (single definition, in order) only until a pass rewrites it.
• Every variant carries a static signature — which CB it moves (produce/consume, n tiles), DST lock effects, arg reads, which thread executes it. CopyTile and its Unpacr expansion carry the same signature, so swapping between levels is sound by construction.
• No catch-all Raw(String) variant. Every line the backend emits is an enum variant; anything opaque is an explicit variant (including inline assembly with a typed operand list).

Pipeline

Fixed order, one pass per concern, each pass TT-specific:

1. Convert + materialize — 1:1 kernel IR → TTOp, then duplicate cross-section values per section (each section is a separate kernel with its own registers and runtime args; arg ordinals stay global; CBs are one hardware object and never duplicated). SSA dedups globally, physical duplicates locally.
2. Sync insertion — ReserveBack/WaitFront/PushBack/PopFront per traffic op.
3. Lock insertion — Acquire placed directly at the LICM position (preheader of the outermost loop containing a pack): the lowering derives the same placement kernel LICM would compute, because the TTIR stream mirrors kernel loop structure 1:1. No effect ops enter the kernel IR for this — the placement is derivable, so there is nothing to search over.
4. Reconfig/init insertion — naive full reconfig before every copy/pack; correctness first, redundancy is safe.
5. Hoist + dedup — adjacent same-config reconfigs collapse; effect ops hoist out of constant-trip loops with trip-count accounting (FIFO depth × trip). Runs on the fully-placed stream: place first, then clean.
6. Regalloc + coalesce — late; linear scan over the linear stream, producing in-place forms.
7. Verify — structural: budgets per id class, CB balance and traffic correspondence, lock pairing, section termination, and the materialization invariant (every value consumed in a section is defined in that section).
8. Render — a table walk, not synthesis.

Each pass rebuilds the stream (drain → push), states which phase it expects (all-virtual → mixed → all-physical), and fails loudly at the exact op it cannot place. Verify guards the composition: whatever the passes did, the final stream must be launchable.

Why no effect ops in the kernel IR

Hoisting the math lock out of the tile loop looks like a job for kernel LICM. It isn’t: the acquire is operand-free, so LICM would trivially hoist it maximally every time — the result is a fixed, derivable position, and the lowering can place the acquire exactly there at insertion time. Adding an effect op to the SSA IR would instead mean auditing CSE (which must never merge two identical acquires straddling a release), DCE, and every restructuring pass — new machinery in the shared IR for zero searchable freedom. Effects whose placement genuinely interacts with state (FIFO trips × loop trip counts, format configs with adjacency) are expressed in TTIR, where the hardware semantics live.

What the physical IR is not for

• No SSA machinery of any kind (subject test).
• No fused-composite matching: at instruction level the chain is the fusion — exp/add/reciprocal sequences share LREGs inside one walk, and the register allocator subsumes the containment analysis. There is no dedicated elu op and none is needed: the decomposition is the only spelling, and it is the fast path.
• No scalar optimization: kernel IR folds all scalar SSA work before lowering; scalar variants render dumbly.
• No restructuring: restructuring is a search dimension, and search lives in the kernel IR.

Migration

The high-level variants are pure abbreviations — macros over lower-level variants with identical signatures. A variant like BinaryTileAdd emits today’s output; when its decomposition into raw instructions is proven, the variant is deleted and nothing is lost. Rust’s exhaustive matches make deletion a compile event that walks through every site, so a half-removed variant cannot exist. The differential test does the proving: the pre- and post-deletion streams must render identically for the whole test corpus. The variant set is the migration ledger — the distance to “no LLK” is countable at any moment by listing the remaining high-level variants.

Other backends

Each backend gets its own physical IR if and when it needs one. The differences between machines are structural (SIMT warps and shared-memory barriers vs FIFOs and DST locks), so a shared physical IR would degenerate into a lowest-common-denominator enum that expresses neither honestly. What is shared: the kernel IR above the boundary, the backend-facing input contract (linear SSA, epilogue-cleaned, arg ordering), and the methodology (typed variants, signature table, phase invariants, rebuild-per-pass, verify, render, differential gate). If two backend physical IRs later converge on the same variants for a machine-independent concept, a shared sub-IR can be extracted then — convergence-discovered unification is sound; designed-upfront unification is a layer nobody asked for.

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.

The Key Insight

In most frameworks, autograd requires a separate graph because the eager execution engine discards intermediate results. Zyx’s tape keeps all graph nodes alive until realize() or drop, and simply prevents their deletion until gradients are computed.

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”

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:

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 nodes 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 values — not the actual data. When the tape is dropped, all tape-preserved nodes 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