The timeline, op by op
A trace holds several planes; the one that matters for kernel work is the device plane (/device:TPU:0), and everything on it has a name you already know how to read. The jit_<function> event is the envelope: wall time per call. Fusions carry their HLO text, so chapter 04's reading skill transfers directly. Custom calls carry their target. Hardware counter samples ride along in the same plane; skip them when reading time, mine them when chasing throttling.
Async copies appear as start and done pairs, and the reading that matters is overlap. In the captured trace, copy-start.1 shows 89.6 µs of in-flight time per iteration, yet the envelope equals the two compute ops' sum exactly: the copy hid completely behind the 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 →, which is chapter 06's pipeline promise showing up as measured fact. When an envelope exceeds the compute sum, the difference is your unhidden transfer time, named and measurable.
# read the device plane programmatically: durations per op, no UI needed
space = xplane_pb2.XSpace()
space.ParseFromString(open(path, "rb").read())
for plane in space.planes:
if "TPU" in plane.name and "Host" not in plane.name:
names = {m.id: m.name for m in plane.event_metadata.values()}
for line in plane.lines:
for ev in line.events:
op = names.get(ev.metadata_id)
us = ev.duration_ps / 1e6 One operational note that cost this site an afternoon: the xplane proto has lived in several packages over the years, and current Colab images ship none of them importable. The lab's setup cell searches the environment, then fetches the proto from the openxla source and compiles it on the spot. Tooling archaeology is part of profiling; budget for it.
How a capture works
There is no magic in a capture, and knowing its shape strips the tool of mystery. Host-side, the runtime records python and dispatch events; device-side, the TPU reports timed op events and counter samples. Everything lands in one container, the XSpace: planes for each device and host, lines within a plane for each stream, and timed events on each line carrying metadata (the op's name, its HLO text) and stats. That is the whole format. GYM·08 came from about forty lines of reading it.
Capture discipline, learned here the measured way: compile outside the window, or the trace is mostly compilation. Trace steady state, enough iterations to swamp noise, and keep the window small, since host-side tracing has real overhead. One file lands per host. And parse programmatically when you want numbers rather than pictures: the viewer and your script read the same bytes.
Check yourself
01 Which plane answers kernel questions, and what is on it?
The device plane (/device:TPU:0): the ops that actually ran, under names that match the compiled module, fusions and custom calls included.
02 Where does a capture come from?
Two recorders: the host runtime logs python and dispatch events while the TPU side logs device execution, and the viewer aligns the two.
Readings
- JAX profiling ↗ where the planes come from
- XLA tools ↗ the dump the event names trace back to