Catching Silent Errors in the Compiler Behind TPU Workloads

When faster code changes the answer

Machine-learning compilers translate model computations into instructions for specialized hardware. Google's XLA:TPU compiler accelerates these computations partly through fusion: combining several operations into one kernel and processing tensors in small tiles. Keeping intermediate results in fast on-chip memory reduces expensive trips to off-chip memory.

That optimization also makes correctness harder to check. A mistake in indexing, memory coordination or numerical conversion can produce incorrect outputs without a crash. Ignacio Cano and colleagues at Google describe a framework for detecting and diagnosing these failures in their preprint, XLA from the Trenches: Correctness on TPU Workloads at Scale.

The authors report that the framework has driven the resolution of more than 142 compiler bugs, with 103 additional findings under review. These are reported production-engineering results, not a proof that the compiler is free of remaining defects.

Turning performance searches into tests

The approach combines synthetic programs with computations drawn from real workloads. Synthetic generation constructs programs with compatible types and tensor shapes, allowing engineers to explore combinations that hand-written tests might miss. However, the authors note that small generated programs often fail to reproduce the memory pressure and operation patterns of production models.

Their complementary strategy uses an existing performance autotuner. This system already tries alternative tile configurations to find faster implementations. Adding output comparisons turns those searches into correctness checks: each candidate's results are compared with the compiler's default configuration within strict numerical tolerances.

The authors argue that this reuses existing infrastructure at near-zero marginal infrastructure cost. It also tests configurations that are not currently defaults but could become defaults after future tuning changes. A defect found outside today's usual execution path may therefore matter to tomorrow's workloads.

AI agents, with human gates

The researchers also describe teams of AI agents that inspect compiler implementations, generate targeted tests, reproduce failures and draft fixes. Workers primarily use instruction-set simulators, reserving physical TPUs for final confirmation.

A central safeguard is checking whether a generated program is valid in the first place. An apparent disagreement can come from a test that violates an operation's semantic requirements rather than from faulty compilation. For problems discovered in low-level compiler passes, agents must also produce a high-level reproducer demonstrating that normal compilation can reach the defective pattern.

Proposed patches receive independent adversarial agent review and human review. Maintainer approval remains necessary before changes land. The authors report that steering agents toward silent numerical corruption, rather than easy-to-detect crashes, helped focus scarce reviewer attention on higher-impact findings.

Finding the first wrong operation

Detection alone does not identify the cause. The framework first extracts standalone subgraphs from a large computation, preserving relevant layouts and configurations. It compares their TPU outputs with unfused execution and a software reference evaluator.

Test inputs require care. Random values can create negative square-root inputs, zero divisors or overflowing reductions, generating misleading discrepancies. The framework works backward through operations to derive safe input bounds. Its reference evaluator also models hardware-specific numerical behavior, including approximations for some mathematical functions and the flushing of very small numbers to zero.

Once a failing fusion is identified, debugging moves inside the intact kernel. Splitting a fusion to inspect intermediate values can change memory use and register allocation, accidentally removing the bug. Instead, the system steps hardware execution and reads intermediate vector-register values, comparing them with expected results.

A rounding error that debugging could hide

One example involves converting a 32-bit value to a narrower format and back, then subtracting the rounded value from the original. This should preserve a rounding residual. The authors describe a lowering behavior that skipped the conversions inside registers, turning the subtraction into a value minus itself and producing zero.

Breaking the fusion apart restored the missing rounding through memory storage, concealing the problem. Inspecting the intact fusion exposed it. The authors report that enforcing explicit rounding in the affected fusion restored the residuals exactly against the reference.

What the evidence establishes

The reported fixes span low-level code generation, algebraic rewrites, precision handling and specialized lowering paths. The largest reported share, 43%, concerns TensorCore low-level optimization and instruction emission.

The scope remains Google's XLA:TPU production environment. Nondeterministic failures can sometimes surface, but they are not the framework's primary target. Its practical lesson is continuous testing coupled with trustworthy comparisons and precise diagnosis—not exhaustive verification or autonomous replacement of compiler maintainers.