The floors

Fusion decides what stays together in fast memory. It says nothing about what a chip actually runs. Once the HLO pipeline finishes, every fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → in the module still has to become real code, and that is the point where XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → stops behaving like one compiler and splits into three.

Below fusion, every backend goes its own way.

On CPU, a fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → becomes LLVM IR wrapped in a thunk runtime (xla/backends/cpu/runtime/thunk.h): a typed sequence of operation nodes an executor fires in dataflow order on a host thread pool. On GPU, a fusion becomes MLIR, or Triton, or a direct call into cuDNN, replayed as its own thunk sequence (xla/backends/gpu/runtime/thunk.h). On TPU, the fusion crosses into libtpu, and libtpu is closed: XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → hands off a program and nothing in this track can read what happens after that hand-off. This is exactly where the kernel path's MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → chapter picks up.

Thunk (xla/backends/cpu/runtime/thunk.h:70) is CPU's uniform dispatch unit: an abstract base class whose Execute is the only virtual call the runtime makes, and whose Kind enum enumerates every concrete flavor a fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → can become, a kernel, a convolution, a dot, a collective, a custom call, and several more. Each thunk declares up front which buffers it reads and writes, and the ThunkExecutor reads those declarations to build a dependency graph, then fires whichever thunk has every input ready, on an Eigen thread pool.

For a small enough computation, eight thunks or fewer, or every buffer under 512 bytes, the executor skips the thread pool and runs the sequence straight through. Dataflow scheduling has a real cost, and paying it to schedule a handful of scalar ops would lose more time than it saves.

GPU fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → resolves through a single dispatcher, GetFusionEmitter, sitting behind a decision tree that runs ahead of it. HloFusionAnalysis looks at each fusion first: if a backend config was already stamped kTritonGemmFusionKind or kCuDnnFusionKind by an earlier pass, the fusion goes straight to Triton or cuDNN. Otherwise the dispatcher inspects the fusion's dominant instruction, its shape decides whether the result is a reduction kernel, a transpose kernel, a scatter, or the generic elementwise loop every unclaimed fusion falls through to.

Triton and cuDNN are not the only way a dot instruction leaves HLO. GemmRewriter can rewrite a supported matmul directly into a __cublas_̲_cublas" style="color:#cc0000">gemm or __cublaslt$matmul custom call, no fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → emitter involved at all, and cuBLAS owns the kernel from there. Which path a given dot takes, Triton-GEMM if it is enabled, cuBLAS otherwise, the generic loop emitter if neither wants it, is settled by a fixed priority order before any of this chapter's codegen machinery runs.

Chapter 1 named PJRT the seam between every frontend and XLA. StreamExecutor (xla/stream_executor/stream_executor.h) is the seam on the other side: one abstract interface per physical device that owns memory allocation, execution streams, kernel loading, and the BLAS and DNN library handles a backend exposes. GPU thunks call through this interface and nothing lower. Stream is the ordered queue thunks enqueue work onto: operations on one stream run in order, operations on different streams can overlap, and a CudaStream (or its ROCm counterpart) is the concrete object underneath.

The description a StreamExecutor returns for its own device, GetDeviceDescription(), is what makes GPU codegen decisions possible at all: threads per warp32 GPU threads scheduled as one unit; the GPU hides latency by switching among resident warps rather than by pipelining a scratchpad.taught in /l/tpu →, shared memory per block, register limits, the exact compute-capability variant. A GPU pass asking whether it can use a particular warp-level trick is really asking this object a question, not guessing from a chip name string.

Custom calls are how a fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → decision gets skipped on purpose. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s FFI (xla/ffi/ffi.h) is a stable, versioned C ABI that lets a custom-call instruction dispatch straight to a handler XLA never compiled: a hand-written kernel, a cuDNN graph, or the pattern-matched kernel the timeline x-ray's online-softmax call reveals in the kernel path's TPU dumps, dispatched by XLA:TPU recognizing that exact program shape.

The call frame a handler receives carries typed buffer descriptors, a sorted attribute list, an execution-stage flag, and a context object for state that needs to survive across calls, all defined by one struct-of-function-pointers vtable. A correct custom call means someone wrote the kernel entirely outside the compiler; the FFI's job is making that handoff safe, not making the kernel unnecessary.