Why is my XLA CPU model only using one core?

Almost always because the executor decided, when the module was compiled, that concurrency was not worth paying for. Three separate conditions can make that call and they are ORed together, so any one of them is enough. The execution graph may already be a chain, meaning every node has an in edge from the node immediately before it, in which case there is no parallelism to find. Every buffer slice that every thunk declares may be at or under 512 bytes, the default execute_sequential_buffer_threshold, which is a per-slice test rather than a total footprint. Or the sequence may hold eight thunks or fewer, the default execute_sequential_num_thunks_threshold. Neither threshold is reachable from an environment variable; both are fields on a ThunkExecutor::Options struct that only a C++ caller can set. So the usual fix for a small model on one core is a larger batch or a larger program rather than a flag. Turn on VLOG(2) for this file and the constructor prints the thunk count, the source and sink node counts, the final is_sequential and small_buffers separately, which tells you which of the three rules fired.

How does the XLA CPU backend decide which thunks can run in parallel?

Purely from declared memory use. Every thunk implements buffer_uses(), which returns the buffer slices it reads and writes, where a slice is an allocation index, a byte offset and a size. ExecutionGraph::Create walks every pair of thunks in order and asks BufferUse::ReadWriteSet whether they conflict. The rule is asymmetric: a write conflicts with any overlapping read or write, while a read conflicts only with an overlapping write. Read-after-read is free, and that single asymmetry is where all of the parallelism in a compiled module comes from. Two slices only overlap inside the same allocation, so two temporaries living at different offsets of one allocation never conflict. A conflicting pair gets an edge. Then a transitive reduction deletes any edge that other edges already imply, which removes an atomic decrement from every future execution, and sets each node priority to the number of nodes reachable from it.

What is the difference between a scheduling edge and an execution edge in XLA?

An execution edge means the dependent thunk must wait for the dependency to finish. It comes from a buffer conflict, so it exists to prevent a data race. A scheduling edge is weaker: the dependent must merely start after the dependency started, and may overlap with it and finish in any order. It comes from a resource conflict on a collective communicator, and it exists so that every process in a job starts its collectives in the same order, which is what keeps them from deadlocking against each other. The executor exploits the difference directly. When a thunk returns an event that is not available yet, ProcessScheduledOutEdges releases only the scheduling edges immediately, before waiting, so a second collective can start while the first is still on the network. That release creates a hazard, since the successor could now finish while the predecessor is still running, so the first scheduling edge also increments the pending_nodes counter, and the comment insists that increment happens before any counter is decremented.

What happens when an XLA CPU thunk never completes?

Nothing visible, which is the hard part. There is no timeout in the executor, no watchdog and no wait loop. Each thread walks its ready queue until it empties and then returns, so the process goes quiet with the cores at zero rather than spinning. The nodes downstream of the stalled thunk keep counters above zero and never enter a ready queue, so pending_nodes never reaches zero, so the execute event handed back to the caller never becomes available. Independent thunks elsewhere in the graph finish normally. The debug assertion that all pending nodes are complete is attached to that same execute event, so it never fires either, and in a release build it is not compiled at all. The practical diagnosis is VLOG(2) on the executor for the graph shape and the sequential decision, VLOG(6) for the full dependency dump with each thunk name and its dependencies, and then a stack dump to find which thunk never signalled its event.

What are the FIFO, LIFO and priority ready queues in the XLA CPU runtime?

Three interchangeable container classes, chosen once by a ready_queue_type field on the options struct and then fixed as a template parameter, so choosing a queue never costs a virtual call on the hot loop. The default is FIFO, which is a vector plus a head index; popping just advances the index and never erases, which makes a pop about three instructions. LIFO is the same vector popped from the back, which favours the node whose predecessor just wrote the buffers it is about to read, so it is the warmer choice for cache. Priority is a std::priority_queue ordered by the number reached in the transitive reduction, which is the count of nodes reachable from each node, on the reasoning that running whatever unblocks the most work keeps the queue fullest. All three also implement PopHalf, which is what the work splitter uses, and they differ in which half they give away: FIFO posts the newest entries, LIFO posts the oldest, and priority keeps the high-priority nodes and posts the rest.

How does XLA spread CPU work across threads without one task per thunk?

By recursive halving. When the loop pops a node and finds more than eight other thunks already ready, and a task runner exists, it calls SplitReadyQueue. That does not post one task per waiting thunk. It calls PopHalf, which takes half the ready queue away as a new queue, and posts that as a single task which will run the same loop and split again. Forty ready thunks reach four threads in three rounds instead of thirty-nine queue pushes, and the comment gives the reason: recursive splitting produces a more uniform distribution and avoids a long tail of work handled by one thread. The cap is a session lock. Thunk::ExecuteSession::kMaxWorkers is four, and the split loop asks with TryJoin, which returns nothing once the cap is reached. The loop breaks on that refusal and keeps processing locally, which is back pressure with no queue and no waiting. The main thread that started the execution is not counted against the cap.

The CPU Thunk Executor: thunk_executor.cc, Line by Line

A GPU gets a driver queue to submit work to. A CPU gets nothing, so the compiled program arrives as a flat list of runtime actions called thunks, and something has to work out what may run at the same time. This 841-line C++ file is that something, read end to end. Each unit declares which buffers it reads and which it writes, and every dependency in the program is derived from those declarations and nothing else, by a conflict test where read-after-read is free. Then a transitive reduction deletes the edges other edges already imply and hands each node a priority equal to how much it unblocks. At run time it is one atomic counter per node and a ready queue: a thunk finishes, decrements its successors, and any counter that hits zero names work that can start now. We read the two edge kinds and why sharing a collective communicator is a weaker constraint than sharing memory; the three thresholds that quietly make a small program single-threaded; the recursive halving that spreads forty ready thunks over four threads without forty queue pushes; and the completion count over sink nodes that makes success and failure end the same way. The signature exhibit is a real ten-thunk graph you step one loop iteration at a time, with all three ready-queue classes.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The CPU Thunk Executor: thunk_executor.cc, Line by Line — built to be understood, not skimmed.

Questions

Why is my XLA CPU model only using one core?: Almost always because the executor decided, when the module was compiled, that concurrency was not worth paying for. Three separate conditions can make that call and they are ORed together, so any one of them is enough. The execution graph may already be a chain, meaning every node has an in edge from the node immediately before it, in which case there is no parallelism to find. Every buffer slice that every thunk declares may be at or under 512 bytes, the default execute_sequential_buffer_threshold, which is a per-slice test rather than a total footprint. Or the sequence may hold eight thunks or fewer, the default execute_sequential_num_thunks_threshold. Neither threshold is reachable from an environment variable; both are fields on a ThunkExecutor::Options struct that only a C++ caller can set. So the usual fix for a small model on one core is a larger batch or a larger program rather than a flag. Turn on VLOG(2) for this file and the constructor prints the thunk count, the source and sink node counts, the final is_sequential and small_buffers separately, which tells you which of the three rules fired.
How does the XLA CPU backend decide which thunks can run in parallel?: Purely from declared memory use. Every thunk implements buffer_uses(), which returns the buffer slices it reads and writes, where a slice is an allocation index, a byte offset and a size. ExecutionGraph::Create walks every pair of thunks in order and asks BufferUse::ReadWriteSet whether they conflict. The rule is asymmetric: a write conflicts with any overlapping read or write, while a read conflicts only with an overlapping write. Read-after-read is free, and that single asymmetry is where all of the parallelism in a compiled module comes from. Two slices only overlap inside the same allocation, so two temporaries living at different offsets of one allocation never conflict. A conflicting pair gets an edge. Then a transitive reduction deletes any edge that other edges already imply, which removes an atomic decrement from every future execution, and sets each node priority to the number of nodes reachable from it.
What is the difference between a scheduling edge and an execution edge in XLA?: An execution edge means the dependent thunk must wait for the dependency to finish. It comes from a buffer conflict, so it exists to prevent a data race. A scheduling edge is weaker: the dependent must merely start after the dependency started, and may overlap with it and finish in any order. It comes from a resource conflict on a collective communicator, and it exists so that every process in a job starts its collectives in the same order, which is what keeps them from deadlocking against each other. The executor exploits the difference directly. When a thunk returns an event that is not available yet, ProcessScheduledOutEdges releases only the scheduling edges immediately, before waiting, so a second collective can start while the first is still on the network. That release creates a hazard, since the successor could now finish while the predecessor is still running, so the first scheduling edge also increments the pending_nodes counter, and the comment insists that increment happens before any counter is decremented.
What happens when an XLA CPU thunk never completes?: Nothing visible, which is the hard part. There is no timeout in the executor, no watchdog and no wait loop. Each thread walks its ready queue until it empties and then returns, so the process goes quiet with the cores at zero rather than spinning. The nodes downstream of the stalled thunk keep counters above zero and never enter a ready queue, so pending_nodes never reaches zero, so the execute event handed back to the caller never becomes available. Independent thunks elsewhere in the graph finish normally. The debug assertion that all pending nodes are complete is attached to that same execute event, so it never fires either, and in a release build it is not compiled at all. The practical diagnosis is VLOG(2) on the executor for the graph shape and the sequential decision, VLOG(6) for the full dependency dump with each thunk name and its dependencies, and then a stack dump to find which thunk never signalled its event.
What are the FIFO, LIFO and priority ready queues in the XLA CPU runtime?: Three interchangeable container classes, chosen once by a ready_queue_type field on the options struct and then fixed as a template parameter, so choosing a queue never costs a virtual call on the hot loop. The default is FIFO, which is a vector plus a head index; popping just advances the index and never erases, which makes a pop about three instructions. LIFO is the same vector popped from the back, which favours the node whose predecessor just wrote the buffers it is about to read, so it is the warmer choice for cache. Priority is a std::priority_queue ordered by the number reached in the transitive reduction, which is the count of nodes reachable from each node, on the reasoning that running whatever unblocks the most work keeps the queue fullest. All three also implement PopHalf, which is what the work splitter uses, and they differ in which half they give away: FIFO posts the newest entries, LIFO posts the oldest, and priority keeps the high-priority nodes and posts the rest.
How does XLA spread CPU work across threads without one task per thunk?: By recursive halving. When the loop pops a node and finds more than eight other thunks already ready, and a task runner exists, it calls SplitReadyQueue. That does not post one task per waiting thunk. It calls PopHalf, which takes half the ready queue away as a new queue, and posts that as a single task which will run the same loop and split again. Forty ready thunks reach four threads in three rounds instead of thirty-nine queue pushes, and the comment gives the reason: recursive splitting produces a more uniform distribution and avoids a long tail of work handled by one thread. The cap is a session lock. Thunk::ExecuteSession::kMaxWorkers is four, and the split loop asks with TryJoin, which returns nothing once the cap is reached. The loop breaks on that refusal and keeps processing locally, which is back pressure with no queue and no waiting. The main thread that started the execution is not counted against the cap.

The CPU Thunk Executor: thunk_executor.cc, Line by Line

Questions

Related explainers