What is an HLO module in XLA?

It is the compilation unit: the object a framework hands to XLA, the object every pass takes and returns, and the object an executable is compiled from. Concretely, an HloModule owns a vector of unique pointers to HloComputation objects and holds a bare pointer to one of them as the entry computation, the way into the program. It also carries a config (the knobs a compile runs under, held copy-on-write so cloned modules can share one), an optional schedule, an input-output alias config saying which outputs may be written on top of which input buffers, a buffer donor config, and a process-unique integer id that ties a dump file back to the module it came from. Ownership is the thing to hold on to: the computations vector is the only place a computation lives, and every other reference to a computation anywhere in the compiler is a pointer into it.

What is the difference between an HloComputation and an HloInstruction?

A computation is a function; an instruction is one operation inside it. The unusual part is that a computation does not store its signature separately. Its parameters are instructions, kept as a vector of pointers to instructions whose opcode is parameter, so num_parameters() is just that vector’s size. Its return value is also an instruction, the one flagged as the root, which is the line prefixed ROOT in the text form. So the whole function, signature included, is one flat list of instructions plus two pointers into it. An instruction, in turn, holds an opcode (one of 134 in the HLO list), the shape of what it produces, an ordered vector of pointers to its operands, and a maintained list of its users, which is the operand edge read backwards. It also knows its parent computation and can carry a sharding and source-location metadata.

What does minor_to_major mean in an XLA layout?

It is the list that says where each dimension of an array sits in memory, and index 0 of the list is the most minor dimension: the one whose index changes fastest as you walk addresses upward. So {1,0} on a two-dimensional array means dimension 1 changes fastest, which is what C and NumPy call row-major, and it is the default XLA assigns when nothing else asks. The source comment states the default as major-to-minor with dimension 0 major, and the code that builds it fills position i with size-1-i. From that one list all the address arithmetic follows: the stride of the most minor dimension is 1, and each successive dimension’s stride is the running product of the sizes of the dimensions more minor than it. The layout prints in braces after the bounds, so f32[4,3,2]{2,1,0} is a shape and its layout together.

Why does a copy instruction appear in my HLO dump when I never wrote one?

Almost always because two operations disagreed about layout. A layout is part of a shape, not a separate annotation, so two buffers with identical shapes and different layouts are not interchangeable: reading one as the other would read the wrong elements. Individual operations have layout preferences, since a matrix multiply may want its right-hand operand arranged one way and a convolution may want channels last, and layout assignment propagates those preferences outward from the operations that hold them. Wherever two propagated preferences meet and disagree, the pass inserts a copy to reconcile them, and that copy is a real instruction with a real cost. The same thing happens across memory spaces: a buffer offloaded to host memory has the same shape with a different memory space in its layout, and moving it back is again a copy.

What does the XLA HLO verifier check?

It runs as an ordinary module pass that never changes anything, and its checks go in a fixed order. First two module-level ones: the module name must not be empty and the entry computation must not be a fusion computation. Then the structure check, which walks the ownership graph itself (every computation’s parent points at this module, every instruction’s parent at its computation, nothing null) and then, in a separate loop, that every operand lives in the same computation as the instruction using it. Then asynchronous start and done pairs, then channel ids, then whether instruction names were allowed to change. Then two visitors walk every computation: a shape verifier that re-derives each instruction’s output shape from its operands and compares, and an instruction verifier carrying the per-opcode structural rules. After the loop come buffers, the entry computation layout, the schedule if the module has one, the alias and buffer-donor configs, and a couple of narrower checks. Any failure is re-raised with a context string naming the stage, which is how you find the pass that broke the graph.

Why do two instructions in different HLO computations get different names?

Because instruction names are unique per module, not per computation, and the module enforces that on the way in rather than complaining. Every computation added to a module runs each of its instructions’ names through a module-wide name uniquer whose separator is a dot. The uniquer keeps a set of used numeric suffixes per name root: the first instruction called out registers suffix 0 and keeps its bare name, and a second one called out finds 0 taken, takes 1, and becomes out.1. The parser adds every non-entry computation first and the entry computation last, so it is usually the entry’s instruction that gets renamed. The practical consequence is that names in a dump are compiler-generated module-wide identifiers that get renumbered whenever a pass clones something, so a line-by-line diff of two dump stages reports every rename as a change. Diff by structure instead.

HLO Module Anatomy: What XLA Holds While It Compiles

An HLO dump is an object graph printed. This walks that graph in the XLA source at a pinned commit: an HloModule that owns a vector of computations and points at one of them as the entry, an HloComputation whose parameters and return value are both just instructions in its own body, and an HloInstruction that owns nothing beyond an opcode, a shape and a list of pointers out to operands and back from users. Then the two things that sit beside the graph and decide most of what it costs: the shape that carries a layout, where minor_to_major turns an index into an address and a disagreement turns into a copy you never wrote; and the schedule, an optional total order added late to hold peak memory down. It closes on the verifier that runs between all two hundred passes, in the order it runs its checks, and the refusals worth knowing by name.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of HLO Module Anatomy: What XLA Holds While It Compiles — built to be understood, not skimmed.

Questions

What is an HLO module in XLA?: It is the compilation unit: the object a framework hands to XLA, the object every pass takes and returns, and the object an executable is compiled from. Concretely, an HloModule owns a vector of unique pointers to HloComputation objects and holds a bare pointer to one of them as the entry computation, the way into the program. It also carries a config (the knobs a compile runs under, held copy-on-write so cloned modules can share one), an optional schedule, an input-output alias config saying which outputs may be written on top of which input buffers, a buffer donor config, and a process-unique integer id that ties a dump file back to the module it came from. Ownership is the thing to hold on to: the computations vector is the only place a computation lives, and every other reference to a computation anywhere in the compiler is a pointer into it.
What is the difference between an HloComputation and an HloInstruction?: A computation is a function; an instruction is one operation inside it. The unusual part is that a computation does not store its signature separately. Its parameters are instructions, kept as a vector of pointers to instructions whose opcode is parameter, so num_parameters() is just that vector’s size. Its return value is also an instruction, the one flagged as the root, which is the line prefixed ROOT in the text form. So the whole function, signature included, is one flat list of instructions plus two pointers into it. An instruction, in turn, holds an opcode (one of 134 in the HLO list), the shape of what it produces, an ordered vector of pointers to its operands, and a maintained list of its users, which is the operand edge read backwards. It also knows its parent computation and can carry a sharding and source-location metadata.
What does minor_to_major mean in an XLA layout?: It is the list that says where each dimension of an array sits in memory, and index 0 of the list is the most minor dimension: the one whose index changes fastest as you walk addresses upward. So {1,0} on a two-dimensional array means dimension 1 changes fastest, which is what C and NumPy call row-major, and it is the default XLA assigns when nothing else asks. The source comment states the default as major-to-minor with dimension 0 major, and the code that builds it fills position i with size-1-i. From that one list all the address arithmetic follows: the stride of the most minor dimension is 1, and each successive dimension’s stride is the running product of the sizes of the dimensions more minor than it. The layout prints in braces after the bounds, so f32[4,3,2]{2,1,0} is a shape and its layout together.
Why does a copy instruction appear in my HLO dump when I never wrote one?: Almost always because two operations disagreed about layout. A layout is part of a shape, not a separate annotation, so two buffers with identical shapes and different layouts are not interchangeable: reading one as the other would read the wrong elements. Individual operations have layout preferences, since a matrix multiply may want its right-hand operand arranged one way and a convolution may want channels last, and layout assignment propagates those preferences outward from the operations that hold them. Wherever two propagated preferences meet and disagree, the pass inserts a copy to reconcile them, and that copy is a real instruction with a real cost. The same thing happens across memory spaces: a buffer offloaded to host memory has the same shape with a different memory space in its layout, and moving it back is again a copy.
What does the XLA HLO verifier check?: It runs as an ordinary module pass that never changes anything, and its checks go in a fixed order. First two module-level ones: the module name must not be empty and the entry computation must not be a fusion computation. Then the structure check, which walks the ownership graph itself (every computation’s parent points at this module, every instruction’s parent at its computation, nothing null) and then, in a separate loop, that every operand lives in the same computation as the instruction using it. Then asynchronous start and done pairs, then channel ids, then whether instruction names were allowed to change. Then two visitors walk every computation: a shape verifier that re-derives each instruction’s output shape from its operands and compares, and an instruction verifier carrying the per-opcode structural rules. After the loop come buffers, the entry computation layout, the schedule if the module has one, the alias and buffer-donor configs, and a couple of narrower checks. Any failure is re-raised with a context string naming the stage, which is how you find the pass that broke the graph.
Why do two instructions in different HLO computations get different names?: Because instruction names are unique per module, not per computation, and the module enforces that on the way in rather than complaining. Every computation added to a module runs each of its instructions’ names through a module-wide name uniquer whose separator is a dot. The uniquer keeps a set of used numeric suffixes per name root: the first instruction called out registers suffix 0 and keeps its bare name, and a second one called out finds 0 taken, takes 1, and becomes out.1. The parser adds every non-entry computation first and the entry computation last, so it is usually the entry’s instruction that gets renamed. The practical consequence is that names in a dump are compiler-generated module-wide identifiers that get renumbered whenever a pass clones something, so a line-by-line diff of two dump stages reports every rename as a change. Diff by structure instead.

HLO Module Anatomy: What XLA Holds While It Compiles

Questions

Related explainers