Why is there a copy instruction in my XLA HLO dump that I never wrote?

Because two operations disagreed about physical layout and the compiler paid for the disagreement on that edge. A shape in HLO carries a layout, which is the permutation saying which dimension of the array moves fastest in memory, and two buffers with the same element type and the same bounds but different layouts are not interchangeable. Layout assignment collects the layout demands that exist in the module, from entry parameters the caller pinned, from infeed and outfeed, from collectives, from layout-constrained custom calls and from backend hooks that pin what cuDNN and cuBLAS accept, and it propagates each demand outward. Wherever a consumer ends up expecting one ordering and the buffer feeding it was assigned another, AssignLayouts calls CopyOperandIfLayoutsDiffer and a kCopy appears on that edge. The comparison it makes is minor-to-major only, so tiles, element size in bits and memory space play no part in it at this stage. Note that this is a different pass from copy insertion, whose copies exist for lifetime reasons rather than layout ones.

What does minor_to_major mean, and what layout does XLA pick by default?

minor_to_major is one list of dimension numbers, and position 0 of it names the most minor dimension, meaning the one whose index changes fastest as you walk memory upward. The comment on the field puts it as a map from physical dimension numbers to logical dimension numbers, with the first element the most minor physical dimension and the last the most major. So on a two-dimensional array {1,0} means dimension 1 moves fastest, which is a contiguous row, which is what C and NumPy call row-major, and {0,1} makes a column contiguous instead. When nothing in the program has a demand, that is what the buffer gets: the source comment states the default XLA layout as major-to-minor with dimension 0 major, and the loop that builds it writes size - 1 - i into position i, so rank 2 gets {1,0} and rank 4 gets {3,2,1,0}. Three other fields ride on the same Layout object and this pass does not decide them: tiles, element size in bits, and memory space, where 0 is the device default and 5 is host memory.

What is the difference between a buffer layout constraint and an operand layout constraint?

A buffer constraint is about where a value is produced; an operand constraint is about what a consumer wants to read. That asymmetry is the whole reason copies exist. A logical buffer is defined exactly once, so when two buffer constraints disagree one of them simply loses on priority and there is nothing to insert. An operand constraint makes no claim about what the producer emits, so when it disagrees with the buffer feeding it the pass can insert a kCopy on that one edge and both constraints hold. There is a third kind, the computation layout constraint, which fixes all the parameters and the result of one computation together and carries a small bitmask recording whether the result, the parameters or both have been set. All three derive from one base class carrying three flags: mandatory, meaning a lower or equal priority may not overwrite it; dfs, deciding which end of the worklist it goes on; and a numeric priority, where a caller-pinned entry layout enters at 3, a layout the pass invented starts at 0, and a pure fallback sits at -2.

Why does XLA insert a copy before a transpose instead of doing the transpose?

Because the pass is not trying to eliminate the data movement, it is trying to relocate it somewhere with a name and a cost. At a transpose the derivation has an explicit stated goal, which the source comment gives as picking the operand layout that makes the transpose a bitcast, a bitcast being a reinterpretation that moves no data at all. The loop is one line of arithmetic: new_minor_to_major[i] is dimensions[minor_to_major[i]], which by construction satisfies the test in ShapeUtil::TransposeIsBitcast, namely that composing the transpose permutation with the output layout gives the input layout. So the pass asks the operand for that layout. If the graph can supply it, nothing is copied anywhere and the transpose really is free. If the operand is an entry parameter the caller pinned, that request loses on priority and a copy appears on the edge instead. Same total bytes moved, but now it is one visible instruction rather than a transpose you would have assumed was cheap.

Should I pin the entry computation layout for my XLA program?

Usually not, and the reason is that pinning does not remove work, it moves work outside the compiler. Pinning is per parameter rather than all-or-nothing: the entry computation layout enters at the top priority if the caller set any layout at all, and each parameter is only stamped when that particular parameter has a layout set. A pinned parameter is a mandatory constraint that nothing derived can override, so a demand arriving from inside the graph loses and a copy is created. Leave it unset and the demand propagates all the way back to the parameter, no copy is created, and the pass writes its chosen layout into the entry computation layout, which becomes the calling convention the runtime has to satisfy. That is usually the better trade, because the runtime arranges the buffer once while a copy inside the graph runs on every step. Pin the one input that genuinely arrives from somewhere you do not control, and leave the rest alone.

Does layout assignment insert copies, or does the copy insertion pass?

Both insert copies and they solve unrelated problems. Layout assignment inserts every copy that exists because two operations wanted different physical arrangements of the same value; search its engine for kCopy and you find CreateCopyWithNewLayout, which builds a plain unary copy for an array and deep-copies a tuple element by element, skipping any element already in the right arrangement. Copy insertion never mentions layout at all: search its 1690 lines for the word and there are zero hits. Its header lists the three problems it does solve, all of them about lifetimes rather than arrangements: an entry parameter or constant that outlives the entry computation and so cannot share its allocation, values that are simultaneously live and must be held in the same buffer, which is the while-loop in-place-update case, and an entry root whose buffer set has to be unambiguous and distinct. Same opcode in the dump, two different reasons for it.

Layout Assignment: Where the Copy in Your HLO Dump Comes From

A copy instruction shows up in your dump between two lines you wrote next to each other, and the only thing that changed across it is the little list in braces. This reads the pass that put it there, at a pinned commit of openxla/xla. What a layout is, which is one permutation of dimension numbers read from the fastest-moving end. The three kinds of constraint the pass collects, and why exactly one of them can be settled by inserting a copy. The switch statement that answers, for all 134 opcodes, whether an operation may hold one arrangement out and a different one in, and the 89 that may not, which is what lets a constraint travel across most of a graph for free. The deque with two ends, the three propagators hanging off it, and the cap of two rewrites that stops it oscillating. The derivation at a transpose, whose stated goal is to make the transpose move no bytes at all, and the copy that lands one edge above it as a result. Then the priority arithmetic that settles a collision, the row-major default handed to whatever nobody argued about, the tuple rule that sends a constraint further upstream instead, and the moment inside AssignLayouts where the copy is finally created. The signature exhibit runs the real propagation loop over a four-instruction module, one pop at a time, with every layout on screen produced by the same arithmetic the pass uses.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Layout Assignment: Where the Copy in Your HLO Dump Comes From — built to be understood, not skimmed.

Questions

Why is there a copy instruction in my XLA HLO dump that I never wrote?: Because two operations disagreed about physical layout and the compiler paid for the disagreement on that edge. A shape in HLO carries a layout, which is the permutation saying which dimension of the array moves fastest in memory, and two buffers with the same element type and the same bounds but different layouts are not interchangeable. Layout assignment collects the layout demands that exist in the module, from entry parameters the caller pinned, from infeed and outfeed, from collectives, from layout-constrained custom calls and from backend hooks that pin what cuDNN and cuBLAS accept, and it propagates each demand outward. Wherever a consumer ends up expecting one ordering and the buffer feeding it was assigned another, AssignLayouts calls CopyOperandIfLayoutsDiffer and a kCopy appears on that edge. The comparison it makes is minor-to-major only, so tiles, element size in bits and memory space play no part in it at this stage. Note that this is a different pass from copy insertion, whose copies exist for lifetime reasons rather than layout ones.
What does minor_to_major mean, and what layout does XLA pick by default?: minor_to_major is one list of dimension numbers, and position 0 of it names the most minor dimension, meaning the one whose index changes fastest as you walk memory upward. The comment on the field puts it as a map from physical dimension numbers to logical dimension numbers, with the first element the most minor physical dimension and the last the most major. So on a two-dimensional array {1,0} means dimension 1 moves fastest, which is a contiguous row, which is what C and NumPy call row-major, and {0,1} makes a column contiguous instead. When nothing in the program has a demand, that is what the buffer gets: the source comment states the default XLA layout as major-to-minor with dimension 0 major, and the loop that builds it writes size - 1 - i into position i, so rank 2 gets {1,0} and rank 4 gets {3,2,1,0}. Three other fields ride on the same Layout object and this pass does not decide them: tiles, element size in bits, and memory space, where 0 is the device default and 5 is host memory.
What is the difference between a buffer layout constraint and an operand layout constraint?: A buffer constraint is about where a value is produced; an operand constraint is about what a consumer wants to read. That asymmetry is the whole reason copies exist. A logical buffer is defined exactly once, so when two buffer constraints disagree one of them simply loses on priority and there is nothing to insert. An operand constraint makes no claim about what the producer emits, so when it disagrees with the buffer feeding it the pass can insert a kCopy on that one edge and both constraints hold. There is a third kind, the computation layout constraint, which fixes all the parameters and the result of one computation together and carries a small bitmask recording whether the result, the parameters or both have been set. All three derive from one base class carrying three flags: mandatory, meaning a lower or equal priority may not overwrite it; dfs, deciding which end of the worklist it goes on; and a numeric priority, where a caller-pinned entry layout enters at 3, a layout the pass invented starts at 0, and a pure fallback sits at -2.
Why does XLA insert a copy before a transpose instead of doing the transpose?: Because the pass is not trying to eliminate the data movement, it is trying to relocate it somewhere with a name and a cost. At a transpose the derivation has an explicit stated goal, which the source comment gives as picking the operand layout that makes the transpose a bitcast, a bitcast being a reinterpretation that moves no data at all. The loop is one line of arithmetic: new_minor_to_major[i] is dimensions[minor_to_major[i]], which by construction satisfies the test in ShapeUtil::TransposeIsBitcast, namely that composing the transpose permutation with the output layout gives the input layout. So the pass asks the operand for that layout. If the graph can supply it, nothing is copied anywhere and the transpose really is free. If the operand is an entry parameter the caller pinned, that request loses on priority and a copy appears on the edge instead. Same total bytes moved, but now it is one visible instruction rather than a transpose you would have assumed was cheap.
Should I pin the entry computation layout for my XLA program?: Usually not, and the reason is that pinning does not remove work, it moves work outside the compiler. Pinning is per parameter rather than all-or-nothing: the entry computation layout enters at the top priority if the caller set any layout at all, and each parameter is only stamped when that particular parameter has a layout set. A pinned parameter is a mandatory constraint that nothing derived can override, so a demand arriving from inside the graph loses and a copy is created. Leave it unset and the demand propagates all the way back to the parameter, no copy is created, and the pass writes its chosen layout into the entry computation layout, which becomes the calling convention the runtime has to satisfy. That is usually the better trade, because the runtime arranges the buffer once while a copy inside the graph runs on every step. Pin the one input that genuinely arrives from somewhere you do not control, and leave the rest alone.
Does layout assignment insert copies, or does the copy insertion pass?: Both insert copies and they solve unrelated problems. Layout assignment inserts every copy that exists because two operations wanted different physical arrangements of the same value; search its engine for kCopy and you find CreateCopyWithNewLayout, which builds a plain unary copy for an array and deep-copies a tuple element by element, skipping any element already in the right arrangement. Copy insertion never mentions layout at all: search its 1690 lines for the word and there are zero hits. Its header lists the three problems it does solve, all of them about lifetimes rather than arrangements: an entry parameter or constant that outlives the entry computation and so cannot share its allocation, values that are simultaneously live and must be held in the same buffer, which is the while-loop in-place-update case, and an entry root whose buffer set has to be unambiguous and distinct. Same opcode in the dump, two different reasons for it.

Layout Assignment: Where the Copy in Your HLO Dump Comes From

Questions

Related explainers