Compiling and Reading the Output

The pipeline from Chapter 3 has been a diagram until now. This chapter opens it up, because a kernel you can read the output of is a kernel you can debug.

Compiling

let ptx_path = compile_kernel(&kernel, &Target::new(env.capability), false)?;

Three arguments:

• the kernel, which supplies the source text and its id,
• the target, a compute capability such as sm_89,
• force, which recompiles even when a cached result exists.

It returns a path, to a file with a .o extension that contains PTX text — // Generated by LLVM NVPTX Back-End is its first line. Beside it, same name, sit two more: .mlir and .rs. The MLIR is the most useful of the three, and the .rs is the captured source from Chapter 8.

Compilation is cached under TEENYC_CACHE_DIR, keyed by the kernel’s id. The first call shells out to teenyc; later calls with the same id return immediately. This matters when you benchmark: an uncached first run measures the compiler, not the kernel. Chapter 18 comes back to it.

Reading the MLIR

This is the compiler’s own record of what it understood your kernel to mean. It is close enough to the source to check line by line, and specific enough to show you what actually happens.

Here is the real MLIR for elemwise_add_forward — the library kernel that vector_add is a copy of — with f32 and a block size of 128:

%c128_i32 = arith.constant 128 : i32
%0 = tt.get_program_id x : i32
%1 = arith.muli %0, %c128_i32 : i32
%2 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>
%3 = tt.splat %1 : i32 -> tensor<128xi32>
%4 = arith.addi %2, %3 : tensor<128xi32>
%5 = tt.splat %arg3 : i32 -> tensor<128xi32>
%6 = arith.cmpi slt, %4, %5 : tensor<128xi32>
%7 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<128x!tt.ptr<f32>>
%8 = tt.addptr %7, %4 : tensor<128x!tt.ptr<f32>>, tensor<128xi32>
%9 = tt.load %8, %6 : tensor<128x!tt.ptr<f32>>
%10 = tt.splat %arg1 : !tt.ptr<f32> -> tensor<128x!tt.ptr<f32>>
%11 = tt.addptr %10, %4 : tensor<128x!tt.ptr<f32>>, tensor<128xi32>
%12 = tt.load %11, %6 : tensor<128x!tt.ptr<f32>>
%13 = tt.splat %arg2 : !tt.ptr<f32> -> tensor<128x!tt.ptr<f32>>
%14 = tt.addptr %13, %4 : tensor<128x!tt.ptr<f32>>, tensor<128xi32>
%15 = arith.addf %9, %12 : tensor<128xf32>
tt.store %14, %15, %6 {operandSegmentSizes = array<i32: 1, 1, 1>} : tensor<128x!tt.ptr<f32>>
tt.return

From kernels/teeny-kernels/tests/snapshots/test_elemwise_add__elemwise_add_forward_mlir.snap.

Now line it up against what you wrote:

Your Rust	The MLIR
T::program_id(Axis::X)	tt.get_program_id x
pid * BLOCK_SIZE	arith.muli %0, %c128_i32
T::arange(0, BLOCK_SIZE)	tt.make_range {start = 0, end = 128}
+ block_start	tt.splat then arith.addi
offsets.lt(n_elements)	tt.splat then arith.cmpi slt
a_ptr.add_offsets(offsets)	tt.splat then tt.addptr
T::load(..., Some(in_bounds), ...)	tt.load %8, %6
a + b	arith.addf
T::store(...)	tt.store %14, %15, %6

Almost one to one. Four things are worth noticing.

BLOCK_SIZE is gone. It is 128, a constant, everywhere it appears — in the multiply, in the range, in every tensor type. This is what “baked in at compile time” means concretely.

Types carry the shape. tensor<128xi32> is a block of 128 integers; tensor<128x!tt.ptr<f32>> is a block of 128 pointers to f32. The type system is tracking your blocks all the way down, and a mismatch here is a mismatch you would have got as a Rust type error first.

tt.splat is broadcasting. Whenever a scalar meets a block, it is copied across every lane. block_start is one integer in your Rust and a tensor<128xi32> here. Three of these appear because three scalars — block_start, n_elements, and each base pointer — get broadcast.

The mask is an argument to the memory operations. tt.load %8, %6 — address tensor, then mask. Nothing branches. That is the whole implementation of the bounds check from Chapter 7: not a jump, just an operand.

The two functions

The full file has two functions, not one. Your kernel appears under a long mangled name, and beside it sits:

tt.func public @elemwise_add_forward_entry_point(...) {
  tt.call @_RINvCslSnLtkXmXla_85elemwise_add_forward_7ef7...(%arg0, %arg1, %arg2, %arg3)
  tt.return
}

That is the wrapper from Chapter 8, doing its one job: giving the loader a predictable symbol to look up. CudaProgram::try_from_ptx resolves {name}_entry_point and gets a function pointer to the thing that calls your kernel.

Snapshot tests

The output above is not a screenshot. It is a committed test fixture, and the pattern behind it is worth stealing:

let kernel = ElemwiseAddForward::<f32>::new(BLOCK_SIZE);
let target = Target::new(Capability::Sm89);
let ptx_path = PathBuf::from(compile_kernel(&kernel, &target, true)?);
let mlir = std::fs::read_to_string(ptx_path.with_extension("mlir"))?;
assert_debug_snapshot!("elemwise_add_forward_mlir", mlir.trim());

This needs teenyc but not a GPU — it compiles, it does not run. So it is the strongest check available on a machine without a card, and it makes any change to the generated code visible in a diff. If you change a kernel and the snapshot moves in a way you did not expect, something happened that you did not intend.

Note Target::new(Capability::Sm89) — a fixed capability rather than the local device’s, so the target is the same everywhere.

These snapshots are not portable between machines. The MLIR embeds the mangled Rust symbol of your kernel, and that name contains rustc’s crate disambiguator — a hash of the crate’s build environment, not of its source. Check out this tree on a second machine and every MLIR snapshot fails, with a diff whose only difference is Cs<something> against Cs<something else>.

The kernel body in the diff will be byte-identical. If that is all you see, nothing is wrong with your kernel. Normalising the disambiguator before comparing would fix it; today nothing does.

So treat a snapshot diff as a question, not a verdict: look at what actually changed inside it before believing it.

When something is wrong

A rough order to work through:

1. Does it compile? A teenyc failure is about the captured text. Check that you have not used anything outside the DSL — Chapter 3’s second consequence.
2. Does the MLIR match your intent? Count the loads and stores. Check the constants. A missing mask operand on a tt.load is visible immediately.
3. Are the numbers wrong at the edges? Suspect the mask, or the other fill value for a reduction.
4. Are the numbers wrong everywhere? Suspect the index arithmetic, or the argument order at the launch site — nothing checks that the tuple you pass matches the parameters your kernel declares.

The end of Part 2

You can write a kernel, run it, and read what the compiler made of it.

That is the whole mechanism. Everything after this is patterns built on it: how to reduce across a row, how to tile a matrix multiply, how to fuse work into a kernel that has already paid for its loads.

Part 3 starts with the first kernel where programs have to do more than mind their own slice.