Lab 02 — Mini Compiler: Expression Language → PTX

Phase: 03 — GPU Compilers | Difficulty: ⭐⭐⭐⭐☆ | Time: 6–10 hours Language: pure Python (stdlib) | Hardware: none — includes a PTX interpreter; optional real-GPU load at the end

Concept primer: ../WARMUP.md Ch. 2 (compiler stages), Ch. 4–5 (PTX), Ch. 8 (the optimizations you'll add).

Run

python solution.py          # compiles 5 kernels to PTX, runs differential tests
python solution.py --dump   # also prints every kernel's emitted PTX

0. The mission

Build the real thing, small: a compiler for elementwise GPU kernels —

kernel saxpy(a: f32, x: f32*, y: f32*) {
    y[i] = a * x[i] + y[i];
}

— that emits legal PTX (the same shape you annotated in Lab 01: .param ABI, cvta.to.global, mad.lo.s32 index, setp guard, fma.rn.f32), where i is implicitly the global thread index and a hidden n parameter guards the tail.

Because most readers don't have a GPU attached, the lab includes a PTX interpreter (PTXInterp) that executes the instruction subset per-thread with float32 rounding — so every compiled kernel is differentially tested against a plain-Python reference. Compiler people call this differential testing; it's how real codegen is validated (alongside conformance suites), and building the interpreter teaches the ISA semantics a second way.

1. The pipeline (read in this order)

Each stage is ~60 lines. The deep lesson: a compiler is a chain of small, testable transforms — after this lab, "we should build a compiler layer" stops being a scary sentence and becomes a costable one.

2. What the emitted PTX looks like

For c[i] = a[i] * 2.0 + b[i]; you'll get exactly the Lab 01 idioms:

.visible .entry scale2(
    .param .u64 scale2_param_0,      // a
    .param .u64 scale2_param_1,      // b
    .param .u64 scale2_param_2,      // c
    .param .u32 scale2_param_3       // n  (implicit)
){
    ...
    mad.lo.s32 %r4, %r1, %r2, %r3;   // i
    setp.ge.s32 %p1, %r4, %r5;
    @%p1 bra DONE;
    mul.wide.s32 %rd10, %r4, 4;      // byte offset, computed once
    ...
    ld.global.f32 %f1, [%rd11];
    mov.f32 %f2, 0f40000000;         // 2.0 as IEEE-754 hex — PTX float literal form
    ld.global.f32 %f3, [%rd12];
    fma.rn.f32 %f4, %f1, %f2, %f3;   // contraction: (a*2.0)+b -> one FMA
    st.global.f32 [%rd13], %f4;
DONE:
    ret;
}

Note 0f40000000: PTX spells float immediates as IEEE-754 bit patterns — one of a dozen small ABI/syntax facts you only learn by emitting the format for real.

3. Extension ladder (this is the lab)

1. Intrinsics: add min(x, y) / max(x, y) → min.f32/max.f32. Touches every stage once — the perfect first change.
2. Constant folding: a pass over the AST: 2.0 * 4.0 + x → 8.0 + x. Assert emitted instruction count drops.
3. CSE: (a[i]+b[i]) * (a[i]+b[i]) currently loads and adds twice. Hash subtrees, reuse registers. Measure again.
4. relu(a[i] * b[i] + c[i]) — fused elementwise kernel via max + folding: you have just built, in miniature, what fusion compilers (XLA, Inductor) do at scale.
5. Real hardware (optional, any NVIDIA GPU):

   from cuda import cuda  # pip install cuda-python
   # cuLinkCreate/cuModuleLoadDataEx accepts PTX text directly:
   err, mod = cuda.cuModuleLoadData(ptx_text.encode())
   

   Launch with cuLaunchKernel and compare against the interpreter. When the driver accepts your compiler's output, you've closed the loop this phase is about.
6. Second backend (Phase 09 foreshadowing): emit C source from the same AST (a for loop over i). Notice what's shared (front/middle end) and what isn't (ABI, index model) — that boundary is the HAL design problem.

4. Common pitfalls

1. Float literals as decimals — PTX wants 0f<bits> form for exactness; emitting 2.0 works in some toolchains and silently loses precision in others. The provided f32_hex() does it right.
2. Reusing the byte-offset register per operand — compute i*4 once; recomputing it per load is correct but bloats the kernel (and your CSE pass should catch it if you regress).
3. Forgetting the guard — the interpreter intentionally runs n not equal to a block multiple; an unguarded kernel writes out of bounds and the differential test catches the corruption. (Same bug as Phase 02 Lab 01 §1.)
4. f32 rounding in the oracle — the interpreter rounds every op to float32 (struct.pack round-trip); compare with relative tolerance, never == (FMA contraction changes the last ulp — WARMUP Ch. 8).

5. What this lab proves about you

"I built a compiler that emits PTX the CUDA driver loads, with differential tests" — that sentence, backed by this repo, is technical credibility with any compiler or kernels team you will ever lead. And the build-vs-adopt argument (WARMUP Ch. 9–10) is no longer abstract: you know exactly which parts were easy (the front end), which were fiddly (ABI details), and which you didn't even attempt (optimization at scale, multi-arch lowering) — which is precisely the cost model an engineering head needs.