GPUReplay: a 50-KB GPU stack for client ML

Running a neural network on a phone's graphics processor requires 48 megabytes of vendor runtime and a 45,000-line driver, most of which exists to serve interactive graphics rather than machine learning — and takes seconds to start up; GPUReplay (Virginia and Purdue) records what that whole stack does at its narrowest interface, then replays the recording on new inputs from a 50-kilobyte program that never learns what any of it means.

Why this problem is worth solving

The thing being replaced is bigger than most people realise, and it runs in the wrong place.

The technique being borrowed usually has a different purpose.

What people did before, and where it broke

• Ship the vendor stack and use it as intended. The universal approach, and reasonable given that the runtime is proprietary. Where it broke. The application inherits everything: the vulnerabilities, the portability constraints, and a startup cost measured in seconds because a trivial job still requires initialising the whole apparatus.
• Cache the expensive parts of startup. Compiled shaders can be cached, and built network pipelines can be cached. Where it broke. These are point fixes for individual bottlenecks. The paper describes its own approach as the same idea pushed to its limit — caching the entire initialisation outcome, at the lowest software layer.
• Reverse-engineer the runtime and write a lean replacement. The obvious route to a small stack. Where it broke. It means decoding proprietary command formats and shader instruction sets, per vendor and often per generation, which is why almost nobody does it.
• Move the graphics stack into a protected environment. What you would do to get on-device machine learning inside a trusted execution environment. Where it broke. Porting tens of megabytes of proprietary runtime into a small secure world defeats the purpose, since the point of that environment is a small trusted base.
• Record and replay for debugging. A mature technique applied to applications, virtual machines and distributed systems. Where it broke. Nothing, but the goal is wrong here. Debugging demands faithful reproduction including nondeterministic events, which is the hard and expensive part — and reproducing the original run is not what a machine learning application needs.
• The gap. The stack’s complexity exists to serve interactive graphics, which produces unpredictable streams of work. Machine learning produces the same fixed sequence of jobs on new data every time, and nobody had asked what a stack for only that would look like.

How a recording replaces the stack

The idea

The stack is complicated for reasons machine learning does not share. The paper’s diagnosis is precise and worth stating carefully, because it explains why a decade of general-purpose engineering is the wrong fit. Interactive graphics generates work nobody can predict — a fifty-second game demo executes 32,000 separate jobs, and a game may even rewrite its shader source as it runs — so the runtime must generate commands and code just in time. Input sizes and buffer lengths change with what the user does, so memory must be managed dynamically. Concurrent programs draw on their own regions of screen, so jobs must be interleaved at fine intervals with separation maintained. Machine learning inverts all three. One application runs a predetermined algorithm, issuing a small fixed set of jobs repeatedly on different data — a typical network is tens of jobs, five or six per layer — so memory and time can be determined ahead of time, and multiplexing happens at intervals of seconds if at all. The stack is charging for flexibility the workload never exercises. Concluding that the main-processor software, rather than the hardware, is what stands between an accelerator and its work is the same diagnosis as the accelerator-initiated storage work covered elsewhere on this site.

Record at the narrowest interface, and you never need to understand anything above it. This is the structural move, and it is what makes the approach practical rather than a reverse-engineering project. However proprietary the runtime, it reaches the hardware through exactly one place: memory-mapped registers and shared memory, mediated by a driver whose source is open. So instrumenting the driver captures everything — the register writes that start jobs, and memory dumps that happen to contain the proprietary commands and shader binaries as opaque bytes. The recorder and replayer stay, in the authors’ words, oblivious to the semantics of most of what they capture. That obliviousness is the feature: no command format is decoded, no instruction set is understood, and the approach works across two vendors, two graphics APIs and several frameworks without change. It also means a job’s entire binary is captured, including branches that were never taken during recording, so conditional logic inside a job replays correctly on new data. The size reduction is the headline — 48 megabytes of runtime and 45,000 lines of driver become 50 kilobytes bare-metal or 10 inside a trusted environment. Shrinking what has to be trusted down to something auditable, and getting fast startup as a consequence, is the same objective as the minimal root-of-trust work covered elsewhere on this site.

Correctness is redefined, which is what makes determinism affordable. Prior record-and-replay systems target debugging, so they must reproduce the original execution faithfully, nondeterministic events included — and capturing and re-injecting those events is where the cost lives. This system needs something weaker: not that the replay resemble the original run, but that the hardware complete the same sequence of jobs. With that definition, nondeterminism splits into two kinds. Sources that would change the outcome are eliminated at record time rather than reproduced — chiefly by disabling asynchronous job submission, which is essential for graphics and nearly free for compute, costing between 2 and 11 percent. Sources that cannot affect the result are simply tolerated — a polling loop whose iteration count varies with flush timing does not need to iterate the same number of times. What remains is a genuine specialisation trade, and the paper is clear about its edges: a recording is one workload on one graphics family, the workload must not branch between different kinds of job depending on input, and recordings port upward across a family but not downward. Accepting a narrower operating envelope in exchange for fitting on a client device is the same bargain as the on-device inference work covered elsewhere on this site.

Recording and replaying one network

The authors report this holding for all 44 networks shipped with three popular frameworks. They note that some architectures use the word branch for routes in their job graph, which are executed unconditionally and therefore pose no problem.

Why it works

• All the proprietary complexity has to pass through one narrow place. Whatever the runtime does, it reaches the hardware only through memory-mapped registers and shared memory → and that interface is mediated by an open-source driver → so instrumenting the driver captures everything without decoding any of it.
• Machine learning issues a fixed set of jobs, so the work can be determined in advance. Graphics generates unpredictable work, which is why commands and code are produced just in time → whereas a network issues the same jobs on new data every run → so all of that generation can happen once, at development time, and be cached as bytes.
• The replay only has to finish the same jobs, not resemble the same run. Debugging replay must reproduce nondeterministic events faithfully, which is the expensive part → but here the hardware only has to complete the same state transitions → so nondeterminism that cannot change the outcome can be ignored entirely.
• Asynchrony is worth almost nothing for compute. Overlapping submission hides management delay behind other work, which matters when jobs are short and numerous → but a compute job runs far longer than it takes to submit → so disabling it costs 2 to 11 percent and removes a major source of divergence.
• Virtual memory decouples the recording from the machine. The open-source driver builds the page tables, so the replayer can place dumps at whatever physical addresses it has → and map them to the virtual addresses the recording expects → which is what lets one recording run on a different device.

What it costs (trade-offs and scope)

• A recording is one workload, frozen. It captures a fixed sequence of jobs, so the workload cannot branch between different kinds of job based on input, and applications that assemble networks programmatically at runtime cannot use a monolithic recording.
• Inference is sometimes slower. On workloads dominated by main-processor overhead the replayer wins substantially, up to 70 percent lower delay on a small network. On larger networks the advantage diminishes and can reverse — the paper reports delays similar or up to 15 percent higher.
• Memory dumps carry data that cannot be excluded. The recorder cannot always tell a job binary from a scratch buffer, so some dumps are large — 66 megabytes for one network on one platform. These regions are mostly zeros and compress well, but they still have to be loaded at replay time.
• Checkpointing is described as generally infeasible. Saving state during replay so a preempted run can resume costs about four times the runtime for one network, because dumping 51 megabytes of memory dwarfs the 45 millisecond replay it protects.
• Recordings travel within a family and mostly in one direction. Moving from a two-core chip to an eight-core one worked with three patches, though without the core-affinity patch performance was four to eight times lower. The reverse direction is described as difficult, requiring proprietary knowledge to recompile shaders.
• Bugs in the recording stack propagate. A semantic error made by the vendor runtime is captured faithfully and shipped. The authors state plainly that their approach neither mitigates nor worsens this, and that the replayer cannot judge semantic correctness at all.
• The recording environment has to be trusted. The security argument assumes the development machine, its stack and its code signing are not compromised, on the grounds that attacking them requires a sustained campaign.
• Instrumentation is manual, per vendor driver. About a thousand lines per driver, amortised across a family but not across vendors, and the authors found automating it less robust than doing it by hand.
• Determinism holds partly because nothing went wrong. Replays were always correct in testing, attributed both to the design disallowing concurrent kernels and to no hardware errors occurring during the benchmarks.

The evaluation

Check yourself

Sources

The paper is “GPUReplay: A 50-KB GPU Stack for Client ML” by Heejin Park and Felix Xiaozhu Lin, published at ASPLOS 2022, pages 157-170, DOI 10.1145/3503222.3507754. The code is public.

This explainer is written from the full text, read from the authors’ copy on arXiv (2105.05085). Everything above comes from the paper’s own sections 2 through 7: the analysis of why graphics stacks are complex, covering just-in-time job generation with the 32,000-job game demo example, dynamic resource management and fine-grained multiplexing, contrasted with the prescribed jobs and coarse multiplexing of machine learning; the trends exploited, including graphics virtual memory managed by an open-source driver, hardware autonomy that lets a job run for milliseconds without processor intervention, and the measurement that enforcing synchronous job submission costs between 2 and 11 percent; the requirement that a workload issue the same kinds of job regardless of input, with the observation that this holds for all 44 networks shipped with three frameworks and that branching within a job is unaffected because the whole binary is dumped; the section 4 recorder covering the interface knowledge needed, the instrumentation points and their amortisation across six and seventeen chips on two vendor drivers, the timing of memory dumps and the reduction from gigabytes to megabytes, and the removal of idle intervals; the section 5 replayer with its static verification of register access, memory access and memory bounds, its 600 line driver, its handoff and checkpointing behaviour, and its detection of divergence with recovery by re-execution; the section 6 implementations across two vendors, several frameworks and four deployment environments with their respective sizes, and the cross-chip recording experiment with its three patches and performance findings; and the section 7 evaluation covering the table of eliminated vulnerabilities, the threat models, startup and inference measurements, recording sizes and their composition, failure injection, granularity trade-offs, preemption delay, and the statement that checkpointing is generally infeasible.