topos

The guide

Emission

The plan as StableHLO text: what lowers, how a raise turns unfold-and-multiply into a convolution, and how the oracle checks the other side.

The plan as text#

A compiled plan is already a closed, statically shaped tensor function: parameters and inputs in, the declared results out. Plan::emit_stablehlo writes it as text — a StableHLO module an industrial runtime such as XLA compiles and runs. Nothing XLA-shaped links into the crate; the text is the boundary, and emission is a sibling of describe, not a second compiler.

module @topos {
  func.func @main(%arg0: tensor<f64>, %arg1: tensor<f64>, %arg2: tensor<f64>) -> (tensor<f64>) {
    %v3 = stablehlo.multiply %arg0, %arg1 : tensor<f64>
    %v4 = stablehlo.subtract %v3, %arg2 : tensor<f64>
    %v5 = stablehlo.multiply %v4, %v4 : tensor<f64>
    return %v5 : tensor<f64>
  }
}
Plate 1. The README's graph as a StableHLO module: a function of the parameter and the inputs, returning the loss. Emission is describe in another syntax.
The code that produced it
use topos::{Detach, Tape};

let (network, [loss]) = Tape::record(|tape| {
    let w = tape.parameter(0.0_f64);
    let x = tape.input(0.0);
    let y = tape.input(0.0);
    let error = w * x - y;
    [error * error].detach()
});
let plan = network.entry([loss]).lower();
show::hlo(plan.emit_stablehlo().expect("every operation lowers"))

The function's arguments are the sources in recording order, parameters first. Its body is the schedule: one StableHLO op per scheduled node, named %vN after the node's index, so a line of the dump and a line of the module are the same node. Constants become stablehlo.constant; a Sum becomes a reduce with an explicit zero; a Broadcast becomes broadcast_in_dim.

A training step as one function#

Because the gradient is recorded as spec, a training step emits as one function: (parameters, batch) -> (loss, gradients...). The host keeps the update loop as plain payload arithmetic and sends the parameters with every batch. makemore_mlp_emitted.rs trains that way through a resident XLA process and grades the losses against the in-crate plan on the page.

module @topos {
  func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<1x2xf32>) -> (tensor<f32>, tensor<2x2xf32>, tensor<1x2xf32>) {
    %v2 = stablehlo.dot_general %arg1, %arg0, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x2xf32>) -> tensor<1x2xf32>
    %v3 = stablehlo.tanh %v2 : tensor<1x2xf32>
    %v4 = stablehlo.multiply %v3, %v3 : tensor<1x2xf32>
    %v5_seed = stablehlo.constant dense<0.0> : tensor<f32>
    %v5 = stablehlo.reduce(%v4 init: %v5_seed) applies stablehlo.add across dimensions = [0, 1] : (tensor<1x2xf32>, tensor<f32>) -> tensor<f32>
    %v6 = stablehlo.constant dense<1.0> : tensor<f32>
    %v7 = stablehlo.broadcast_in_dim %v6, dims = [] : (tensor<f32>) -> tensor<1x2xf32>
    %v8 = stablehlo.constant dense<0.0> : tensor<1x2xf32>
    %v9 = stablehlo.add %v8, %v7 : tensor<1x2xf32>
    %v10 = stablehlo.multiply %v9, %v3 : tensor<1x2xf32>
    %v11 = stablehlo.multiply %v9, %v3 : tensor<1x2xf32>
    %v12 = stablehlo.constant dense<0.0> : tensor<1x2xf32>
    %v13 = stablehlo.add %v12, %v10 : tensor<1x2xf32>
    %v14 = stablehlo.add %v13, %v11 : tensor<1x2xf32>
    %v15 = stablehlo.constant dense<1.0> : tensor<1x2xf32>
    %v16 = stablehlo.multiply %v3, %v3 : tensor<1x2xf32>
    %v17 = stablehlo.subtract %v15, %v16 : tensor<1x2xf32>
    %v18 = stablehlo.multiply %v14, %v17 : tensor<1x2xf32>
    %v19 = stablehlo.constant dense<0.0> : tensor<1x2xf32>
    %v20 = stablehlo.add %v19, %v18 : tensor<1x2xf32>
    %v23 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<1x2xf32>) -> tensor<2x1xf32>
    %v24 = stablehlo.dot_general %v23, %v20, contracting_dims = [1] x [0] : (tensor<2x1xf32>, tensor<1x2xf32>) -> tensor<2x2xf32>
    %v27 = stablehlo.constant dense<0.0> : tensor<2x2xf32>
    %v28 = stablehlo.add %v27, %v24 : tensor<2x2xf32>
    return %v5, %v28, %v3 : tensor<f32>, tensor<2x2xf32>, tensor<1x2xf32>
  }
}
Plate 2. The plan as StableHLO: the same schedule, written as text an industrial compiler accepts. The function takes the parameter and the input and returns exactly what the entry declared, in that order.

Result order is declared#

The plan's results are the entry's roots then its observes, in the order they were declared, deduplicated to the first occurrence. Emission does not invent an output convention; it returns exactly that list, and Plan::results is the same list as symbols.

entry roots:    5 (the loss), 28 (its gradient)
entry observes: 3 (hidden)
plan.results(): [5, 28, 3]

func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<1x2xf32>) -> (tensor<f32>, tensor<2x2xf32>, tensor<1x2xf32>) {
    return %v5, %v28, %v3 : tensor<f32>, tensor<2x2xf32>, tensor<1x2xf32>
Plate 3. Result order is declared. The plan's results are the entry's roots then its observes, in that order, and the emitted function returns exactly that list.
The code that produced it
use topos::{Tape, Tensor};

let tape: Tape<f32> = Tape::new();
let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
let hidden = input.matmul(weights).tanh();
let loss = (hidden * hidden).sum();
let adjoints = tape.differentiate(loss, [weights]);
let (weights, hidden, loss) = (weights.symbol(), hidden.symbol(), loss.symbol());
let network = tape.into_network();

let plan = network.entry(adjoints.roots()).observe([hidden]).lower();
let module = plan.emit_stablehlo().expect("every operation lowers");
let signature = module.lines().find(|line| line.contains("func.func")).unwrap_or("").trim();
let returned = module.lines().find(|line| line.trim().starts_with("return")).unwrap_or("").trim();
let results: Vec<String> = plan.results().iter().map(|symbol| symbol.index().to_string()).collect();
show::lines([
    format!("entry roots:    {} (the loss), {} (its gradient)", loss.index(), adjoints.of(weights).index()),
    format!("entry observes: {} (hidden)", hidden.index()),
    format!("plan.results(): [{}]", results.join(", ")),
    String::new(),
    signature.to_string(),
    format!("    {returned}"),
])

Raises#

The catalog's matches have a second job. At home an elected pattern calls a fused kernel; abroad it becomes the other compiler's named op. The im2col-and-multiply chain a convolution records becomes stablehlo.convolution; a max-pool window becomes reduce_window; the batch-norm formulas become batch_norm_training and batch_norm_inference. That is how a small instruction set still hands an industrial backend the op it is good at, without a Convolution instruction and its own backward rule.

module @topos {
  func.func @main(%arg0: tensor<2x1x3x3xf32>, %arg1: tensor<2xf32>, %arg2: tensor<1x1x6x6xf32>) -> (tensor<1x2x2x2xf32>) {
    %v7 = stablehlo.transpose %arg0, dims = [1, 2, 3, 0] : (tensor<2x1x3x3xf32>) -> tensor<1x3x3x2xf32>
    %v8 = stablehlo.reshape %v7 : (tensor<1x3x3x2xf32>) -> tensor<9x2xf32>
    %v9_kernel = stablehlo.reshape %v8 : (tensor<9x2xf32>) -> tensor<1x3x3x2xf32>
    %v9_windows = stablehlo.convolution(%arg2, %v9_kernel) dim_numbers = [b, f, 0, 1]x[i, 0, 1, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x1x6x6xf32>, tensor<1x3x3x2xf32>) -> tensor<1x4x4x2xf32>
    %v9 = stablehlo.reshape %v9_windows : (tensor<1x4x4x2xf32>) -> tensor<16x2xf32>
    %v10 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<2xf32>) -> tensor<16x2xf32>
    %v11 = stablehlo.add %v9, %v10 : tensor<16x2xf32>
    %v12 = stablehlo.reshape %v11 : (tensor<16x2xf32>) -> tensor<1x4x4x2xf32>
    %v13 = stablehlo.transpose %v12, dims = [0, 3, 1, 2] : (tensor<1x4x4x2xf32>) -> tensor<1x2x4x4xf32>
    %v14 = stablehlo.constant dense<0.0> : tensor<1x2x4x4xf32>
    %v15 = stablehlo.maximum %v13, %v14 : tensor<1x2x4x4xf32>
    %v27_seed = stablehlo.constant dense<0xFF800000> : tensor<f32>
    %v27 = "stablehlo.reduce_window"(%v15, %v27_seed) ({
    ^bb0(%v27_left: tensor<f32>, %v27_right: tensor<f32>):
      %v27_larger = stablehlo.maximum %v27_left, %v27_right : tensor<f32>
      stablehlo.return %v27_larger : tensor<f32>
    }) {window_dimensions = array<i64: 1, 1, 2, 2>, window_strides = array<i64: 1, 1, 2, 2>} : (tensor<1x2x4x4xf32>, tensor<f32>) -> tensor<1x2x2x2xf32>
    return %v27 : tensor<1x2x2x2xf32>
  }
}
Plate 4. A convolution recorded as pad, unfold, permute, reshape, and a matrix product, emitted as stablehlo.convolution; a max-pool as reduce_window. The same match that fuses at home raises abroad.
The code that produced it
use topos::{Shape, Tape, Tensor, conv2d, init, max_pool};

let tape: Tape<f32> = Tape::new();
let image = tape.input(init::normal(1, 1.0)(&Shape::new([1, 1, 6, 6])));
let filters = tape.parameter(init::normal(2, 0.3)(&Shape::new([2, 1, 3, 3])));
let bias = tape.parameter(Tensor::filled([2], 0.0_f32));
let convolved = conv2d(image, filters, bias, 1, 0).relu();
let pooled = max_pool(convolved, 2, 2).symbol();
let network = tape.into_network();

let module = network.entry([pooled]).lower().emit_stablehlo().expect("every operation lowers");
assert!(module.contains("stablehlo.convolution"));
assert!(module.contains("stablehlo.reduce_window"));
show::hlo(module)

Compositions, written out#

What did not earn an instruction is emitted as the composition it is. One attention head is a product, a scale, an additive mask, the softmax spelled from the stable log-softmax, and a product with the values; the module says so, and the receiving compiler is free to fuse it back.

module @topos {
  func.func @main(%arg0: tensor<4x8xf32>, %arg1: tensor<4x8xf32>, %arg2: tensor<4x8xf32>) -> (tensor<4x8xf32>) {
    %v3 = stablehlo.constant dense<[[0.0, 0xFF800000, 0xFF800000, 0xFF800000], [0.0, 0.0, 0xFF800000, 0xFF800000], [0.0, 0.0, 0.0, 0xFF800000], [0.0, 0.0, 0.0, 0.0]]> : tensor<4x4xf32>
    %v4 = stablehlo.constant dense<0.35355338> : tensor<f32>
    %v5 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<4x8xf32>) -> tensor<8x4xf32>
    %v6 = stablehlo.dot_general %arg0, %v5, contracting_dims = [1] x [0] : (tensor<4x8xf32>, tensor<8x4xf32>) -> tensor<4x4xf32>
    %v7 = stablehlo.broadcast_in_dim %v4, dims = [] : (tensor<f32>) -> tensor<4x4xf32>
    %v8 = stablehlo.multiply %v6, %v7 : tensor<4x4xf32>
    %v9 = stablehlo.add %v8, %v3 : tensor<4x4xf32>
    %v10_low = stablehlo.constant dense<0xFF800000> : tensor<f32>
    %v10_peak = stablehlo.reduce(%v9 init: %v10_low) applies stablehlo.maximum across dimensions = [1] : (tensor<4x4xf32>, tensor<f32>) -> tensor<4xf32>
    %v10_spread_peak = stablehlo.broadcast_in_dim %v10_peak, dims = [0] : (tensor<4xf32>) -> tensor<4x4xf32>
    %v10_centered = stablehlo.subtract %v9, %v10_spread_peak : tensor<4x4xf32>
    %v10_exp = stablehlo.exponential %v10_centered : tensor<4x4xf32>
    %v10_total_seed = stablehlo.constant dense<0.0> : tensor<f32>
    %v10_total = stablehlo.reduce(%v10_exp init: %v10_total_seed) applies stablehlo.add across dimensions = [1] : (tensor<4x4xf32>, tensor<f32>) -> tensor<4xf32>
    %v10_normalizer = stablehlo.log %v10_total : tensor<4xf32>
    %v10_spread_normalizer = stablehlo.broadcast_in_dim %v10_normalizer, dims = [0] : (tensor<4xf32>) -> tensor<4x4xf32>
    %v10 = stablehlo.subtract %v10_centered, %v10_spread_normalizer : tensor<4x4xf32>
    %v11 = stablehlo.exponential %v10 : tensor<4x4xf32>
    %v12 = stablehlo.dot_general %v11, %arg2, contracting_dims = [1] x [0] : (tensor<4x4xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
    return %v12 : tensor<4x8xf32>
  }
}
Plate 5. One attention head as StableHLO: the composition the facade recorded, written out for another compiler. The stable softmax is the only fused core in it.
The code that produced it
use topos::{Shape, Tape, Tensor, causal_mask, init, scaled_dot_product};

let tape: Tape<f32> = Tape::new();
let query = tape.input(init::normal(1, 1.0)(&Shape::new([4, 8])));
let key = tape.input(init::normal(2, 1.0)(&Shape::new([4, 8])));
let value = tape.input(init::normal(3, 1.0)(&Shape::new([4, 8])));
let mask = tape.leaf(causal_mask(4, f32::NEG_INFINITY));
let scale = tape.leaf(Tensor::from(1.0 / 8f32.sqrt()));
let head = scaled_dot_product(query, key, value, mask, scale).symbol();
let network = tape.into_network();
show::hlo(network.entry([head]).lower().emit_stablehlo().expect("every operation lowers"))

Conformance#

The slow run still checks the other side of the boundary. When TOPOS_STABLEHLO_VALIDATOR and TOPOS_STABLEHLO_EVALUATOR point at a toolchain — tools/validate-stablehlo.py and tools/eval-stablehlo.py, which use the MLIR bindings a jax install bundles — every emitted module in the test suite is parsed by an external parser and executed by the StableHLO reference interpreter, and its results are checked against the plan's own within an envelope. Where float physics forbids bit-identity the check is a tolerance, and a miss is reported, not skipped. CI runs this job.

On a language model, the in-crate run, the in-crate plan, and XLA on CPU produce the same text. The same module through a vendor GPU plugin ran faster and was wrong — and the disagreement was provable because three independent implementations agreed with each other and it did not. GPT-2 on topos tells it in full.

Where next#