Knok: Tensor Graphs as Rust Build Artifacts

Knok is a static-shape tensor graph compiler for Rust. You define a graph in build.rs; Knok traces it during cargo build, lowers it to MLIR, compiles it with IREE, and generates a typed Rust wrapper around the result.

If those names are unfamiliar, MLIR is compiler infrastructure for representing and transforming programs at several levels, while IREE is an ML compiler and runtime built on top of it. In Knok, MLIR is the intermediate form where tensor operations can be checked and optimized, and IREE turns that graph into executable code for a CPU or GPU and provides the runtime used to load and run it.

use knok_build::prelude::*;

#[knok_build::graph(backend = Backend::LlvmCpu)]
fn forward(x: T2<f32, 2, 2>) -> T2<f32, 2, 2> {
    relu(matmul(x.clone(), x) + 1.0)
}

fn main() {
    knok_build::compile_graphs!(forward);
}

The generated graph is imported from normal Rust code:

use knok::prelude::*;

knok::generated_graphs!(pub mod graphs);

let y = graphs::forward::call(x)?;

The graph, compiled artifact, backend metadata, and Rust signature are produced together. There is no separate model file whose function names and shapes have to be duplicated in application code.

This is the layer I wanted after working on Eerie, my Rust binding for the IREE runtime. Eerie handles execution, but it assumes that an IREE module already exists. Knok handles the graph before it reaches the runtime.

Why build.rs

Rust projects already use build.rs to generate protocol bindings, compile native libraries, process shaders, and embed assets. Tensor compilation fits the same model: the graph is a source input, the IREE VMFB is a build artifact, and generated Rust code is its interface.

A bad graph now fails during the build. Changing the graph regenerates the artifact. The wrapper carries the input and output types, so shape information is not maintained separately in a config file or a runtime call.

Knok still needs iree-compile, either on PATH or through KNOK_IREE_COMPILE. The repository uses Nix to provide a pinned compiler. Packaging the compiler cleanly for users is still one of the rougher parts of the project.

Replacing the parser with tracing

The first Knok frontend used a procedural macro that parsed graph function bodies. It worked for small examples, but it quickly became a second language disguised as Rust. Helper functions, loops, tuples, local bindings, literals, and expression reuse all needed custom parser support.

Knok 0.3 replaced that frontend with build-time tracing. Graph functions are now ordinary Rust functions operating on traced tensor values. Rust handles the host language; Knok handles tensor operations.

fn block(x: T2<f32, 2, 2>) -> T2<f32, 2, 2> {
    relu(matmul(x.clone(), x) + 1.0)
}

#[knok_build::graph(backend = Backend::LlvmCpu)]
fn stacked(mut x: T2<f32, 2, 2>) -> T2<f32, 2, 2> {
    for _ in 0..4 {
        x = block(x);
    }
    x
}

The loop runs on the build host and appends four blocks to the graph. Knok does not need its own implementation of Rust loops or helper functions.

Control flow depending on tensor values is different. Those values do not exist while build.rs is running, so a normal Rust if cannot branch on them. Data-dependent behavior has to use graph operations such as where. That boundary is much clearer than trying to teach a macro an increasingly large subset of Rust.

Static shapes

Knok stores tensor dimensions in the type:

Tensor1<f32, 128>
Tensor2<f32, 32, 64>
Tensor4<f32, 1, 224, 224, 3>

This is restrictive, but deliberate. I am mostly interested in graphs used in robotics, control, and embedded inference. Those interfaces are usually fixed anyway: a policy has a fixed observation size, a controller has a fixed state vector, and a perception model is normally deployed at a chosen resolution.

For those cases, static shapes make the generated API useful. A Tensor2<f32, 3, 4> cannot be passed where Tensor2<f32, 4, 3> is expected, and the mismatch is caught before execution.

Knok currently supports ranks zero through six and a practical set of arithmetic, reductions, layout operations, indexing, matrix multiplication, convolution, and pooling operations.

Backends

Knok currently exposes four IREE backends:

• Backend::LlvmCpu compiles through IREE’s LLVM CPU backend and runs with the local-task driver.
• Backend::MetalSpirv targets Metal on macOS.
• Backend::VulkanSpirv targets Vulkan and is enabled with the vulkan feature.
• Backend::Cuda targets NVIDIA GPUs and is enabled with the cuda feature.

The backend is selected on the graph:

#[knok_build::graph(backend = Backend::Cuda)]
fn forward(x: T2<f32, 32, 32>) -> T2<f32, 32, 32> {
    relu(matmul(x.clone(), x) + 1.0)
}

The generated artifact records the matching runtime driver. Engine::for_artifact uses that metadata, so application code does not have to duplicate the CPU, Metal, Vulkan, or CUDA choice at runtime.

For repeated execution, it is better to create an engine once:

let engine =
    knok::Engine::for_artifact(graphs::forward::artifact())?;

for x in inputs {
    let y = graphs::forward::run(&engine, x)?;
}

This keeps the IREE instance, device, loaded module, and other runtime state alive between calls.

Why kernel fusion matters

Consider a small graph:

relu(matmul(x, w) + bias)

A naive eager implementation may launch a matrix multiplication kernel, write an intermediate tensor, launch an addition kernel, write another tensor, and then launch ReLU. The arithmetic in the last two operations is cheap. Launching kernels and moving intermediate tensors through memory may cost more than the operations themselves.

A compiler can see the whole graph and fuse compatible operations. The matrix multiplication may remain a specialized kernel, while bias addition and ReLU are folded into its output stage. That removes intermediate allocations, cuts memory traffic, and reduces launch overhead.

Fusion matters even more for chains of small elementwise operations:

relu(clamp(x * scale + bias, -1.0, 1.0))

Running every operation separately is a poor fit for GPUs, where each launch has a fixed cost and global memory access is expensive. A fused kernel can load the input once, perform the arithmetic in registers, and write the final result once.

This is one reason Knok builds a graph instead of behaving like an eager tensor library. Ahead-of-time compilation is not only about producing code for another backend. It gives IREE enough context to remove intermediates, fuse kernels, specialize static shapes, and schedule the graph for the selected hardware.

Autograd

Knok also has experimental reverse-mode autodiff. A graph returning a scalar loss can be passed to grad_graphs!:

#[knok_build::graph(backend = Backend::LlvmCpu)]
fn loss(x: T1<f32, 4>) -> T0<f32> {
    mean(square(x))
}

fn main() {
    knok_build::grad_graphs!(loss);
}

The generated wrapper returns the value and input gradient:

let (loss, dx) =
    graphs::loss_value_and_grad::call(x)?;

The differentiation pass works on Knok’s graph IR before MLIR lowering. It produces another graph, which goes through the same compiler and runtime path as the original one. Autodiff stays out of both the tracing frontend and the runtime.

It is not a complete training framework. There are no optimizers, parameter containers, data loaders, or eager execution. The current goal is a differentiable static graph compiler.

Examples

knok-demo is an interactive egui application that runs fixed 1024 x 1024 graphs for Mandelbrot rendering, heat diffusion, wave simulation, Conway’s Game of Life, and particle interaction. It includes a plain ndarray CPU implementation for comparison and can build Knok variants for LLVM CPU, Metal, Vulkan, and CUDA. It is mainly a demonstration that compiled graphs can sit inside an ordinary Rust GUI application rather than a standalone model runner.

knok-mnist-training dogfoods the autograd path by training a fixed-shape MNIST MLP. Cargo builds one graph for inference and another for the scalar loss and parameter gradients. The surrounding training loop is normal Rust: it parses the dataset, shuffles batches, applies SGD updates, and calls the compiled graphs. There is also an egui binary for training the model and testing mouse-drawn digits. It is a small example, but it shows the intended split: Knok compiles the numerical graph while the rest of the application remains ordinary Rust code.

Current limitations

Compiler distribution and caching need more work. Operator and gradient coverage are incomplete. Dynamic shapes are outside the current typed frontend. External MLIR files can be compiled through the same wrapper-generation path, but there is no general ONNX or PyTorch importer yet.

GPU support also needs broader testing across machines and IREE compiler builds.

Knok is not intended to replace PyTorch or JAX. It is for static tensor programs that are deployed as part of a Rust system and should be built, typed, and packaged with that system.

Eerie handles the IREE runtime. Knok adds the build-time graph layer above it.