Passing a TPU Through to a Paravirtualized Guest

A customer asked us to run Google TPUs inside our zones. The obvious way to hand a guest a piece of PCI hardware is IOMMU-backed passthrough, and Xen has supported that for HVM guests for years.

Our zones default to PV, without hardware virtualization extensions or a VT-x or AMD-V dependency. That is a big part of why they boot and fork as fast as they do. Moving these workloads to a different guest model would have meant giving up properties we deliberately built the runtime around, so the interesting problem was making passthrough work without changing the execution model.

Xen's PV mode had never had an IOMMU story for a domU. Every DMA a PV guest issues goes out under dom0's address space, because PV was designed before device passthrough to an unprivileged domain was something anyone needed from it.

The TPU work ended up being three problems stacked on top of each other. We had to give Xen an IOMMU interface a PV guest could actually use, teach Linux's driver stack to sit on top of that interface instead of a real IOMMU, and fix everything in our own toolstack that quietly assumed passthrough only happens to HVM guests.

All three shipped. NVIDIA GPU passthrough on PV came from the same work because none of the resulting interface was TPU-specific.

Why Xen Had No IOMMU For domU

Xen's IOMMU code has always been dom0-only in spirit even where it wasn't in name. Devices get a context implicitly by being assigned to a domain, and nothing asks a guest to manage its own mappings because that was never part of the model.

There is no hypercall for a domU to say "map this guest frame to this device," because a PV guest doesn't have the address space assumptions that would make that request coherent without an IOMMU already in the picture translating for it.

We didn't start from nothing here. Teddy Astie had already posted an RFC upstream reworking the IOMMU subsystem around per-domain IOMMU contexts and exposing them through a hypercall called PV-IOMMU.

The redesign makes contexts first-class objects a domain can allocate, populate, and attach devices to instead of giving each device one implicit context tied to its owning domain. That is what a guest managing its own IOMMU mappings actually requires. The existing code had nowhere for a guest-driven request to land. We ported that work forward and adapted it for a PV guest.

Parts of Xen's IOMMU interface had changed since the RFC was posted, and the original interface was written around dom0 assumptions. A guest also sees PCI devices differently from the host, so the interface needed to account for a guest managing mappings for devices it already owns. Making that work on real hardware exposed another set of problems.

GCE's IOMMU is, in Alex's words, "particularly... particular," and getting AMD-Vi and VT-d to behave under a paravirtualized guest surfaced a long list of latent bugs. Guest-driven contexts exercise paths that had not been exercised this way before. We found completion handling problems, incorrect invalidation behavior, unnecessarily expensive IOTLB flushing, and cases where guest-controlled inputs needed tighter bounds.

These were not specific to PV-IOMMU. They were AMD-Vi and VT-d paths that a dom0-only, implicitly assigned IOMMU had never exercised in quite the same way.

Teaching the Linux Kernel to Run as PV-IOMMU Client

None of the Xen work matters if the guest kernel has nothing to talk to it with.

The Linux changes teach the kernel to run as a PV-IOMMU client. That means giving a guest kernel a way to communicate with the hypervisor-managed IOMMU and making Linux understand that the IOMMU it is talking to is not necessarily physical hardware owned by the guest.

Two of the fixes come from assumptions baked into Linux's PCI and IOMMU code that change when the IOMMU is mediated by the hypervisor.

The iommu-dma reserves PCI MMIO windows out of a device's DMA address space by default. That is correct when there is a real IOMMU translating and wrong when the translation model is being provided differently. Under paravirtualized DMA, those reserved ranges can become address space the guest actually needs.

MSI-X has a similar issue. A PV guest cannot behave as though it owns the interrupt controller in the same way a hardware virtualized guest can. Xen already owns interrupt routing for the PV domain, so PCI interrupt handling has to respect that model.

There is also plumbing needed to make device passthrough usable from xen-pcifront. The guest and hypervisor have different views of PCI identity, and IOMMU operations ultimately have to refer to the physical device the hypervisor knows about. Mapping operations also need to be efficient enough that large workloads are not paying the cost of one hypercall for every individual page.

What Our Own Toolstack Got Wrong

After the Xen and Linux changes were in place, the PCI passthrough path in our own stack still did not work. The remaining problems had little to do with IOMMU internals and much more to do with code that assumed HVM.

PV domains were not being created as domains capable of owning IOMMU-backed devices. That assumption had never mattered before because PCI devices had never been assigned to a PV guest.

The interrupt wiring also assumed the device was still associated with dom0 at a point where it had already been assigned to the guest. That made the failure show up much later than the actual mistake.

The fix was to wire the interrupt in the guest's own interrupt space instead of trying to route it through a host-side device that no longer owned it. Even with the device correctly assigned and interrupts wired, DMA was still silently wrong.

A PV guest cannot simply hand a device its own view of physical memory and assume those addresses mean the same thing to the hardware. Something has to translate the guest's page view into the actual machine frames the device can access.

Without that translation, DMA targets the wrong physical memory. The guest can ask a device to perform DMA using an address that looks valid from inside the guest while the device is actually pointed somewhere completely different in machine memory.

The translation path existed, but enabling it had become tied to an unrelated device configuration setting. That coupling had never mattered when PV guests did not own passthrough devices.

We separated those concerns so DMA translation is enabled because passthrough requires it, not because an unrelated device setting happened to turn it on.

Where The Time Was Actually Going

Once it worked, it needed to work fast enough for a real libtpu workload, which allocates large amounts of memory and expects mapping and unmapping it to be cheap.

Unmapping in particular showed latency that initially looked like an IOMMU problem.

The timing data pointed instead at scheduler noise from asynchronous page-table updates. The guest issues an unmap, the update to the page tables happens on a different schedule than the guest expects, and the guest ends up waiting on interrupts introduced by that asynchronous path rather than on any IOMMU hardware limit.

We added tracing around the page-table update path to make that visible instead of guessing at it, and left speeding up the asynchronous path itself as follow-up work.

The correctness fix was worth shipping on its own. The performance work is separate and still needs to be finished.

Shipping It Across Xen, Linux, and the Runtime

‍

The implementation spans Xen, the guest kernel, and our own runtime. All three layers had to change together because device passthrough crosses all three.

There is still more work to do around the higher-level Kubernetes device experience and performance, but the underlying PCI passthrough path is no longer tied to a specific accelerator.

By the time the work was done, very little of it was actually TPU-specific. A PV guest with a correctly assigned IOMMU context, correctly wired interrupts, and correctly translated DMA does not care what is on the other end of the PCI bus. NVIDIA GPU passthrough on PV started working as a side effect of the same set of fixes.

The TPU was the forcing function. The PV-IOMMU interface is the reusable part.