notebook № 236 plates · 25 min to read, longer to play ·
Four ways to program a GPU
On how much to spell out
part 5 of 6 in How chips compute · by Slavik Shynkarenko
You can tell a GPU what one lane does, what one thread does, what one tile does or what happens to the whole array. Each step up hands more decisions to the compiler. This notebook walks the four levels with the same few programs, so you can see what you write, what gets decided for you and what it costs.
plate I · the unit
Who counts the lanes?
In the GPU notebook I took the machine apart: warps, schedulers, tiles of shared memory. This one sits on the other side of the table. When you sit down to program that machine, the first thing you choose, often without noticing, is the unit you write code for.
There are four common answers. You can write for one lane, spelling out a vector instruction that works on eight floats at once. You can write for one thread and let the hardware run thousands of copies. You can write for one tile, a block of data, and let a compiler decide which thread touches which element. Or you can write for the whole array and let a library decide everything. Each step up, the compiler takes over a decision you used to make.
Fig. 23.1 runs the same program in each: SAXPY, y = a·x + y, the hello-world of parallel code. The interesting lines aren’t the maths. They’re the ones about the machine: how wide a step is, which element is mine, and what to do with the leftovers when n doesn’t divide evenly. The last tab is my own language, Miri, which I’ll come back to in plate VI.
Set n to 1,000, then nudge it to 1,003. Which versions need code for the three extra elements, and which just mask them off?
plate II · simd and simt
Who owns the width
SIMD (single instruction, multiple data) is what CPUs offer: registers that hold several numbers, and instructions that work on all of them at once. An SSE register holds four floats, AVX2 eight, AVX-512 sixteen, Arm’s NEON four. The width is written into the code. The instruction names carry it, the loop steps by it, and the tail is sized by it. Move to a wider unit and you rewrite the kernel, or keep a version per machine.
A GPU makes the opposite deal. You write SPMD code, single program, multiple data: the program describes one thread and says nothing about width. The hardware groups the threads itself: 32 to a warp on NVIDIA, 64 to a wavefront on AMD’s datacentre chips and 32 or 64 on its gaming ones, 32 to a SIMD-group on Apple, and 8 to 32 on Intel, where the compiler chooses. The same per-thread code fits every one of those widths, and the divergence rules from the GPU notebook are the price of that convenience.
The idea isn’t tied to GPUs. Intel’s ispc compiles SPMD programs to CPU vector units, and a good auto-vectoriser does a quieter version of the same trick on ordinary loops. SIMT is SIMD with the width moved out of your code and into the machine.
Switch the CPU from SSE to AVX-512 and count the changed lines. Then switch the GPU from NVIDIA to AMD CDNA and count again.
plate III · subgroups
Passing notes in the warp
The thread model pretends each thread is alone. Sooner or later they have to combine their results, and the classic way is to talk through memory. To sum 32 numbers, each thread writes its value to shared memory, everyone waits at a barrier, half the threads add a neighbour’s value, everyone waits again, and so on down to one.
Warp-level operations skip the detour. __shfl_down_sync lets a thread read a register belonging to another thread of the same warp, with no shared memory and no barrier, because the warp already runs in step. Five shuffles sum 32 values. Recent NVIDIA GPUs go one further and reduce a whole warp of integers in a single instruction.
Every API has these now, under different names: WaveActiveSum in HLSL, subgroupAdd in Vulkan’s GLSL, simd_sum in Metal. The portable word for the group is subgroup. It’s the one place the thread model admits the group exists, and the fastest kernels lean on it hard.
Step through both versions and compare the counters. Which of them would hurt most once the block has 1,024 threads instead of 32?
plate IV · tiles
Write for the tile
Look back at the shared-memory sum in plate III, or the tiled matrix multiply in the GPU notebook. Much of that code is bookkeeping: which thread loads which element, when to wait, how to fetch the next tile while computing this one. In 2019 Philippe Tillet’s Triton made a simple bet: let the programmer write for a tile and leave the threads to the compiler.
In Triton you write a program instance, a function that owns one block of the output and works on whole tiles at once with tl.load, tl.dot and tl.store. The compiler decides how many warps run it, how the tile splits across their threads, what goes into shared memory and when to use tensor cores. You choose the tile shape, and the tile shape decides almost everything else.
Fig. 23.4 shows why that one choice matters. A bigger tile reuses each byte it loads more times, so its arithmetic intensity (flops per byte moved) rises and the kernel climbs the roofline from memory-bound towards compute-bound. But the tile has to fit in on-chip memory, double-buffered, and on a small matrix a few huge tiles leave most of the GPU idle. Tile size is a three-way trade between reuse, on-chip memory and parallelism, and Triton ships an autotuner to search it for you.
The idea spread fast. Triton now writes the GPU code behind PyTorch’s compiler, JAX has Pallas, NVIDIA has added a tile model to CUDA itself, and CUTLASS’s CuTe gives C++ programmers tile abstractions with full control underneath.
Pick the 1,024 matrix and push both tile sides to 256. The intensity looks great. Why is the kernel slow anyway?
plate V · arrays
Write for the whole array
At the top of the ladder you stop writing kernels at all. y = layer_norm(gelu(x @ w + b) + r) in PyTorch or JAX is a sentence about whole arrays, and the library picks the kernels, the block sizes and the widths. This is how most GPU code gets written today, by people who rarely think about a warp, and that’s fine.
The cost hides between the operations. Run eagerly, each operation is its own kernel, and each kernel reads its input from GPU memory and writes its output back. The matrix multiply is compute-bound and earns its trip. The bias, the GELU, the residual add and the norm each make a full round trip for a sliver of arithmetic. Most of their time goes to moving bytes, and on small batches to launching kernels at all.
Fusion is the fix. A compiler like XLA or torch.compile sees the whole expression, merges the cheap operations into one kernel that keeps values in registers, and folds some of them into the matrix multiply itself. torch.compile does it by writing Triton kernels, the tile model of plate IV, behind your back. The array level gets its speed from the levels below it.
Switch to 64 rows, the shape of a single decoding step in an LLM. What eats the time now, and how much does fusion buy back?
plate VI · choosing
Picking a level
None of the levels wins outright. Lower levels give you control and ask for code, expertise and a version per machine. Higher levels give you short programs that run everywhere and ask you to trust a compiler. Most real projects mix them: an array program for the model, a tile kernel for the one operation that matters, a library call for the matrix multiply.
My rule of thumb: start at the highest level that can express the job, and drop a level only when a profiler tells you to. Portability points the same way. Array libraries and portable APIs like Vulkan and WebGPU run on every vendor’s GPU, CUDA on one.
Miri is my attempt at a level the map was missing. In Miri you mark data gpu to say where it lives and write the per-element body as a forall loop in the ordinary program. The language, not you, decides how that loop becomes kernels, launches and data movement. I wanted the ease of the array level with the per-element control of a kernel, and above all I wanted what you write kept apart from how the machine runs it. A loop says nothing about warps or tiles, so the same loop can be lowered to a thread-level kernel on one target and a tile-level program on another, which is what a language built for many back ends needs.
gpu let a = [1.0, 2.0, 3.0, 4.0]
gpu let b = [5.0, 6.0, 7.0, 8.0]
gpu var dst = [0.0, 0.0, 0.0, 0.0]
forall i in 0..a.length()
dst[i] = a[i] + b[i]Every one of these programs still has to become something a GPU can execute, and that trip is its own story: PTX, SPIR-V, Metal’s intermediate code and the compilers hiding in your graphics driver. It’s the subject of the next notebook.
Pick “one kernel, every vendor”. Which levels survive, and what does each give up to get there?
Questions people ask
- What is the difference between SIMD and SIMT?
SIMD code names the vector width: an AVX2 instruction adds eight floats at once, and a wider unit needs different instructions. SIMT code is written for one thread, and the hardware runs threads in groups (32 on NVIDIA, 32 or 64 on AMD) without the code saying so. Underneath, both are lanes running in lockstep.
- What is tile-based GPU programming?
You write a program for one block of the data, such as a 128 × 128 tile of a matrix, using whole-tile operations like load, dot and store. The compiler decides how threads split the tile, what goes into shared memory and when to use tensor cores. Triton and JAX’s Pallas work this way.
- Is Triton easier than CUDA?
For most deep-learning kernels, yes. A Triton matrix multiply is a few dozen lines of Python, against hundreds of lines of tuned CUDA, and it gets close to vendor libraries on common shapes. CUDA still wins when you need control the tile model hides, such as unusual memory layouts or warp-level tricks.
Related material
elsewhere in the notebooks
further reading
- CUDA C++ Programming Guide NVIDIA The SIMT model and the warp-level functions of plate III, from the source.
- Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations Tillet, Kung and Cox, 2019 The paper behind plate IV, and its argument that blocks, not threads, are the right unit for neural-network kernels.
- Triton tutorials Triton Vector add, fused softmax, matrix multiply and fused attention, each a short program you can run.
- Pallas: a JAX kernel language JAX The same tile idea inside JAX, for GPUs and TPUs.
- Intel Implicit SPMD Program Compiler (ispc) Intel Per-thread programs compiled to CPU vector units: the bridge between plates I and II.
- Vulkan Subgroup Tutorial Khronos Subgroup operations from the portable side, with diagrams of what each one does.
- Making Deep Learning Go Brrrr From First Principles Horace He Compute-bound, memory-bound, overhead-bound: the three regimes behind plates IV and V.
- miri-lang.org My language: data marked gpu, per-element work written as forall loops, and the compiler deciding how they run.
- Why I’m building a GPU-first programming language Why I think the GPU should be the default target of a language, not the exotic one.