Codegen and the backend seam

There is a fast way to see where the shared part of XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → ends, and it is a directory listing. xla/hlo/transforms holds the hardware-independent passes: the simplifiers, the expanders, the collective transforms, host offloading, precision propagation. Everything past that point lives under xla/backends, and xla/backends has exactly five entries: autotuner, cpu, gpu, interpreter, profiler.

No tpu. That absence is not an oversight and it is the shape of this whole page. The TPU backend is not in the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → tree, so the boundary between what you can read and what you cannot is not somewhere fuzzy inside codegen. It is a directory that does not exist.

The official architecture describes three stages, and the five directories map onto the last two. First, target-independent optimization: common subexpression elimination, target-independent operation 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 →, buffer analysis. Second, backend-specific HLO-level optimization done with target-specific information and needs in mind, including pattern matching for library calls. Third, code generation, with the CPU and GPU backends leveraging LLVM for low-level IR, optimization, and code generation. The first stage is shared. The second and third are per backend, and one backend keeps them in a wheel.

xla/hlo/transforms/          # hardware-independent, shared by every backend
  collectives/
  expanders/
  simplifiers/
  host_offloader.h  bfloat16_propagation.h  defuser.h  ...

xla/backends/                # everything past the shared pipeline
  autotuner/
  cpu/
  gpu/
  interpreter/
  profiler/
                             # there is no tpu/ here, and that is the point

The chapter above walks the emitters, and the thing worth adding here is that each of these backends produces two artifacts, not one, and confusing them makes the runtime look more mysterious than it is. LLVM generates the machine code for a kernel. A separate structure, the thunk sequence, is the program that decides which kernel runs when and against which buffers. Compiled code and execution plan, produced by the same compile, consumed by different parts of the runtime.

GPU adds a complication the CPU path does not have: not every operation becomes XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →-generated code at all. A matmul can be rewritten into a cuBLAS custom call, 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 be handed to Triton, and a convolution can go to cuDNN. Each of those is a decision made before codegen about whether XLA emits anything for that instruction, and the fusion kind kCustom from the fusion unit's lesson is the marker for a fusion claimed by that kind of path.

The interpreter backend in that listing is easy to skip past and worth a sentence. It is the reference implementation, the thing run_hlo_module compares a compiled result against by default. A backend that produces a different answer from the interpreter has a bug, and having a slow, obvious, hardware-free implementation in the tree is what makes that statement checkable.

libtpu is one file that holds three things people usually think of as separate. Its own publisher describes it as the core library that enables JAX, PyTorch, and TensorFlow to execute models on Google Cloud TPUs, providing core functionality for compilation, inter-chip communication, and runtime execution. So: a compiler, a network layer, and a runtime, shipped as a wheel. Version 0.0.45 published on 2026-08-01, at 215.9 MB.

The fourth thing in there is the one the pjrt lesson established. It is a PJRT plugin, exporting GetPjrtApi, discovered by JAX through the same mechanism as any other plugin. Every layer above it, IFRT, jax.jit, your model code, is talking to a function-pointer struct and does not know or care that a compiler is on the other side.

What stays readable through that boundary is more than you would expect, and /xla/pipeline is built on it. The closed compiler dumps through the same flags, so its pass names come out: Phase_1_pre_layout_assignment_passes, X64_elimination, hlo_device_type_async_wrapper, add-random-host-offloading, tpu-embedding-thread-annotator. You cannot read the implementation of a single one of those. You can read the whole list, in order, and diff the module on either side of any of them.

A closed compiler still prints its own pass names, and still dumps the module on both sides of each one.

Now the question this lesson exists for. If you are building a backend, what do you have to write? The contract is narrower than the size of the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → tree suggests, and the reason is in a struct from the pjrt lesson. PJRT_Client_Compile takes a PJRT_Program, and a PJRT_Program is a code blob plus a format string. Nothing in that contract says the code is HLO. Nothing says XLA's pass pipeline runs.

So 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 optimizer is a library available to a backend, not a stage a backend has to pass through. You can take the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → you were handed, lower it into your own IR, and never construct an HloModule at all. Or you can link XLA, run its pipeline, and write only an emitter, which is what the CPU and GPU backends do. Both choices satisfy the same PJRT contract, and a framework above cannot tell which one you made.

There is a conformance harness for the interface itself. RegisterPjRtCApiTestFactory gives a plugin the standard tests of basic PJRT C API behaviours, and the integration guide's own acceptance criteria are ordinary JAX operations: element-wise arithmetic, a jit compile, and a pmap with a collective reduction. That last one is the real gate, because it exercises the multi-device path that the key-value callbacks in the pjrt lesson exist to bootstrap.

Both routes exist in public, at both extremes of scale, which makes the decision concrete rather than theoretical. The example plugin under xla/pjrt/plugin/example_plugin is the linked route at its smallest: a C++ PjRtClient subclass, wrapped by pjrt::CreateWrapperClient, exported through a PJRT_Api that pjrt::CreatePjrtApi assembled. libtpu is the other route at its largest: an entire compiler and runtime behind the same exported symbol, sharing none of 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 codegen.

The question that decides which one fits is not about performance, it is about how close your hardware's execution model is to the one 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 passes assume. If shapes, layouts, and buffers mean roughly what they mean on a GPU, linking the pipeline saves you years. If they do not, the pipeline will spend those years fighting you, and the PJRT contract is deliberately loose enough to let you skip it.

Either way the surface a framework sees is identical, and that is the property the ifrt unit's lesson depends on. A layer above PJRT can treat every backend as the same kind of object precisely because the seam was drawn at a struct of function pointers rather than at an IR.