What does buffer assignment do in XLA?

It decides which bytes of device memory every intermediate value of a compiled program lives in. It runs after all HLO rewriting and after the instructions have been put in a total order, and it does not change the program: it produces a separate BufferAssignment object that maps each value to an allocation, a byte offset and a size, and the code generator queries that object while emitting kernels. The result is that a compiled XLA program never allocates at run time. It asks the runtime for a small list of slabs, one for each entry parameter, one for the result, usually one large one for all the temporaries, and every kernel already has its offsets baked in. The pass runs from the GPU and CPU compilers through the same function, BufferAssigner::CreateAssignment, which builds alias analysis, collects the schedule, computes live ranges, colours the values, splits computations into global and thread-local, and then assigns.

How does XLA decide two tensors can share the same memory?

By comparing live ranges. Once the module has a schedule, every value gets a start and an end as integer positions in that schedule: the start is where the value is written, and the end is the last instruction that reads it. Two values may share an address when those intervals do not overlap. The intervals are closed, so a value written at instruction seven and a value last read at instruction seven do overlap, because both are alive at seven. There is one exception, and it is where in-place updates come from: when one range ends exactly where the other starts, the pass asks whether the later value is produced by an instruction that reads the earlier one and writes its output on top. If it is, the two do not interfere and one buffer serves both. Every other refusal is logged in words, which is why running with a raised verbosity on this file tells you exactly why two buffers you expected to share did not.

What is fragmentation in an XLA buffer assignment report?

It is the gap between the memory the program was assigned and the memory it strictly needed. The floor is computed by a heap algorithm that does no packing at all: walk the schedule, add the sizes of everything live at each point, take the largest total. Fragmentation is the assigned total minus that floor, and it is only computed when every computation has a schedule, since otherwise there is no floor to compare against. Two things create the gap. Alignment rounds every buffer start forward, which on the GPU means a 256-byte boundary because cuBLAS requires it, so a 1000-byte buffer effectively occupies 1024. And best fit is a heuristic: it can leave a hole that nothing later happens to fill. The number matters because device memory is a hard wall, and on a model sized to use most of it a few percent of fragmentation is the difference between running and an out-of-memory error.

What is the difference between mandatory aliasing and buffer reuse in XLA?

Mandatory aliasing is a fact about the program; reuse is a decision the compiler made. Alias analysis groups values into HloBuffers, and two values in one HloBuffer must occupy the same memory for the program to be correct. A while loop is the clearest case: the initial value, the loop instruction output, the body parameter, the body root and the condition parameter are five different values and one buffer. In-place operations are the case that shows up in ordinary models, since a dynamic-update-slice, a scatter and an in-place all-reduce start alias their operand with their output by definition. Reuse is different: two values at the same offset in one allocation merely happen to share, because the packing proved their lifetimes do not overlap and it saved bytes. Buffer assignment cannot break your program by reusing wrongly, because the mandatory grouping was fixed before it started.

Why did a copy instruction appear near my dynamic-update-slice?

Because something still reads the value the update writes over. A dynamic-update-slice aliases its first operand with its output, so the output is the input at the same address with a few elements changed. That is only safe if nobody needs the old contents afterwards. A pass called copy insertion handles this conservatively: it first adds a copy of the operand for every in-place operation in the module, then a second phase reruns alias analysis and deletes each copy it can prove is unnecessary, iterating to a fixed point. When the update is the only remaining reader, the copy disappears and you never knew it existed. When a later instruction reads the original, the proof fails, the copy stays, and it costs a full extra tensor plus the kernel that fills it. Removing that later read is usually what makes the copy go away.

Which files in an XLA dump show memory usage?

Four, all written alongside the module text at the point where a buffer assignment exists. The one ending -buffer-assignment.txt lists every allocation with its size and flags, and under each one the values assigned to it with their sizes and offsets; two values at the same offset means their lifetimes were proven disjoint. The one ending -memory-usage-report.txt answers a different question, what is large: it groups values by offset, sorts by size, and prints a running cumulative total with percentages and the shapes involved, stopping early once the rest is under five percent. Two more appear only if you pass the buffer-assignment-analysis dump flag: -buffer-assignment-values.txt lists every used value, and -live-range.txt prints the flattened instruction sequence with indices, every value as a start and end in those indices, and the set of values live at the peak moment, which is the direct answer to why the peak is what it is.

Buffer Assignment: How Every Value Gets an Address

A compiled XLA program never allocates. It receives a handful of slabs from the runtime and every kernel was compiled with its offsets inside those slabs already baked in. This walks the pass that picks them, at a pinned commit of openxla/xla: the values that must share memory because the language says so, the schedule that turns a dependency graph into a clock so that a live range can be two integers, the closed-interval interference test and the one endpoint case it makes an exception for, the colour that is a memory space rather than a graph colouring, the four kinds of value that never reach the packing, and the biggest-first best-fit heap that packs everything left. Then the nine ordered reasons an allocation refuses a buffer, the in-place update whose copy you never see until something reads the old value, and the four files a dump writes with the answer in them. The signature exhibit packs a real schedule with the real rules and lets you drag a live range until the sharing breaks.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Buffer Assignment: How Every Value Gets an Address — built to be understood, not skimmed.

Questions

What does buffer assignment do in XLA?: It decides which bytes of device memory every intermediate value of a compiled program lives in. It runs after all HLO rewriting and after the instructions have been put in a total order, and it does not change the program: it produces a separate BufferAssignment object that maps each value to an allocation, a byte offset and a size, and the code generator queries that object while emitting kernels. The result is that a compiled XLA program never allocates at run time. It asks the runtime for a small list of slabs, one for each entry parameter, one for the result, usually one large one for all the temporaries, and every kernel already has its offsets baked in. The pass runs from the GPU and CPU compilers through the same function, BufferAssigner::CreateAssignment, which builds alias analysis, collects the schedule, computes live ranges, colours the values, splits computations into global and thread-local, and then assigns.
How does XLA decide two tensors can share the same memory?: By comparing live ranges. Once the module has a schedule, every value gets a start and an end as integer positions in that schedule: the start is where the value is written, and the end is the last instruction that reads it. Two values may share an address when those intervals do not overlap. The intervals are closed, so a value written at instruction seven and a value last read at instruction seven do overlap, because both are alive at seven. There is one exception, and it is where in-place updates come from: when one range ends exactly where the other starts, the pass asks whether the later value is produced by an instruction that reads the earlier one and writes its output on top. If it is, the two do not interfere and one buffer serves both. Every other refusal is logged in words, which is why running with a raised verbosity on this file tells you exactly why two buffers you expected to share did not.
What is fragmentation in an XLA buffer assignment report?: It is the gap between the memory the program was assigned and the memory it strictly needed. The floor is computed by a heap algorithm that does no packing at all: walk the schedule, add the sizes of everything live at each point, take the largest total. Fragmentation is the assigned total minus that floor, and it is only computed when every computation has a schedule, since otherwise there is no floor to compare against. Two things create the gap. Alignment rounds every buffer start forward, which on the GPU means a 256-byte boundary because cuBLAS requires it, so a 1000-byte buffer effectively occupies 1024. And best fit is a heuristic: it can leave a hole that nothing later happens to fill. The number matters because device memory is a hard wall, and on a model sized to use most of it a few percent of fragmentation is the difference between running and an out-of-memory error.
What is the difference between mandatory aliasing and buffer reuse in XLA?: Mandatory aliasing is a fact about the program; reuse is a decision the compiler made. Alias analysis groups values into HloBuffers, and two values in one HloBuffer must occupy the same memory for the program to be correct. A while loop is the clearest case: the initial value, the loop instruction output, the body parameter, the body root and the condition parameter are five different values and one buffer. In-place operations are the case that shows up in ordinary models, since a dynamic-update-slice, a scatter and an in-place all-reduce start alias their operand with their output by definition. Reuse is different: two values at the same offset in one allocation merely happen to share, because the packing proved their lifetimes do not overlap and it saved bytes. Buffer assignment cannot break your program by reusing wrongly, because the mandatory grouping was fixed before it started.
Why did a copy instruction appear near my dynamic-update-slice?: Because something still reads the value the update writes over. A dynamic-update-slice aliases its first operand with its output, so the output is the input at the same address with a few elements changed. That is only safe if nobody needs the old contents afterwards. A pass called copy insertion handles this conservatively: it first adds a copy of the operand for every in-place operation in the module, then a second phase reruns alias analysis and deletes each copy it can prove is unnecessary, iterating to a fixed point. When the update is the only remaining reader, the copy disappears and you never knew it existed. When a later instruction reads the original, the proof fails, the copy stays, and it costs a full extra tensor plus the kernel that fills it. Removing that later read is usually what makes the copy go away.
Which files in an XLA dump show memory usage?: Four, all written alongside the module text at the point where a buffer assignment exists. The one ending -buffer-assignment.txt lists every allocation with its size and flags, and under each one the values assigned to it with their sizes and offsets; two values at the same offset means their lifetimes were proven disjoint. The one ending -memory-usage-report.txt answers a different question, what is large: it groups values by offset, sorts by size, and prints a running cumulative total with percentages and the shapes involved, stopping early once the rest is under five percent. Two more appear only if you pass the buffer-assignment-analysis dump flag: -buffer-assignment-values.txt lists every used value, and -live-range.txt prints the flattened instruction sequence with indices, every value as a start and end in those indices, and the set of values live at the peak moment, which is the direct answer to why the peak is what it is.

Buffer Assignment: How Every Value Gets an Address

Questions

Related explainers