Intermediate Examples¶
The intermediate examples add reduction loops, normalization, and transformer-oriented elementwise operations. Each script is a runnable single-device Golden Harness case.
All five scripts declare a2a3, a2a3sim, a5, and a5sim. Repository CI
exercises a2a3, a2a3sim, and a5sim.
GEMM¶
This example extends the beginner matmul by tiling K. The first K tile creates
the accumulator with pl.matmul; later tiles update it with
pl.matmul_acc.
Key topics: M/N/K blocking, sequential reduction, and FP32 accumulation.
LayerNorm¶
The hidden dimension fits in one tile. Each row tile computes the mean and variance, normalizes its input, then applies gamma and beta with broadcast operations.
Key topics: row_sum, row and column broadcasts, reshaping reduction results,
and numerical tolerances.
RMSNorm¶
RMSNorm demonstrates a reduction whose hidden dimension is larger than one tile. One pass accumulates the sum of squares across hidden chunks; a second pass normalizes each chunk and applies gamma.
Key topics: chunked reductions, persistent accumulators, rsqrt, and a
two-pass kernel structure.
RoPE¶
The RoPE example splits each head into two halves and applies the rotary position transform with column broadcasts.
Key topics: transformer tensor layout, half-vector slicing, broadcast multiplication, and assembling an output from slices.
Softmax¶
This is a numerically stable row-wise softmax: subtract the row maximum, exponentiate, reduce the denominator, and broadcast the division.
Key topics: stable reductions, row_max, row_sum, and row broadcast.
Suggested reading order¶
Use this order when learning the DSL:
- GEMM for a loop-carried tile accumulator.
- Softmax for a compact reduction-and-broadcast pipeline.
- LayerNorm for multiple dependent reductions.
- RMSNorm for a reduction split across hidden chunks.
- RoPE for transformer-specific slicing and layout.