Four kinds, named in the enum
The chapter above makes the load-bearing argument, that fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → changes where intermediates live and never changes the math, and it is the right thing to hold onto. What a chapter does not have room for is the vocabulary, and the vocabulary is small enough to learn in one sitting. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → distinguishes four kinds of fusion, and it prints the kind on every fusion instruction it emits.
kLoop is the default shape: the fused computation's root is the primary node, and codegen emits an element-at-a-time loop over it. kInput also has the primary node at the root but generates code close to what the unfused version would produce, which is the shape reductions take. kOutput is the one where the primary node is not the root, and it carries a constraint the other kinds do not: one operand buffer has to alias the output buffer. kCustom is the backend's own category for anything that fits none of these.
In this repo's TPU capture, every fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → in all three programs prints kind=kLoop, including the ones that exist only to wrap a single bitcast. That uniformity is itself informative: on this backend and these programs, the interesting variation is not in the kind field, it is in what got pulled inside the computation.
enum class FusionKind {
kLoop,
kInput,
kOutput,
kCustom,
};The queue, and the formula that orders it
Fusion on GPU runs as a priority queue, and the priority has an actual formula written in the header. PriorityFusion is an HloModulePass whose name() returns priority-fusion, and its documented algorithm is: compute priority = time_unfused - time_fused for each producer, put the producers with positive benefit into a queue ordered by that benefit, then pop the top one, fuse it, and update the priorities of everything the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → touched.
Read the formula rather than the word and one common mistake disappears. The queue is ordered by estimated time saved, not by how many operations a merge absorbs. A fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → that swallows one enormous intermediate outranks a fusion that swallows six small ones, because the difference of two time estimates is dominated by the memory traffic the merge removed.
The queue is ordered by time saved, not by operations absorbed.
The private members of the pass name the machinery behind the estimate. GpuHloCostAnalysis::Options cost_analysis_options_ is described in the header as what defines the priorities in the queue, and HloFusionAnalysisCache fusion_analysis_cache_ keeps that analysis from being recomputed for every candidate as the graph changes underneath it. There is a third member worth knowing about as a debugging surface: std::unique_ptr<FusionProcessDumpProto> fusion_process_dump_, which logs the decisions the queue made. Fusion decisions are dumpable as their own artifact, separately from the HLO text.
Reading fusion names in a dump
A fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → appears twice in a module and the two appearances look confusingly similar. There is an instruction, printed inside a computation, that looks like %fusion.1 = bf16[1024] fusion(%copy-done, %reduce_max.7), kind=kLoop, calls=%fused_computation.2. And there is the computation itself, printed at module scope with its own parameter list and its own ROOT. The calls= field is the only link between them.
The numbers on the two are independent counters, which is the trap worth meeting deliberately. In the softmax capture, the ROOT instruction %fusion calls %fused_computation.1, and the instruction feeding it, %fusion.1, calls %fused_computation.2. Matching a fusion.42 to a fused_computation.42 by number will be wrong more often than right. Follow the calls= field.
Once you are inside the computation, the parameter list tells you what the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → actually reads, which is more useful than the instruction's operand list because it is annotated with shapes. %fused_computation.2 (param_0.9: bf16[1024,512], param_1.10: bf16[1024]) says this fusion consumes the whole matrix and a vector of row maxima. The ROOT tells you what it produces and, in this case, what anchored it: ROOT %reduce.1 = bf16[1024] reduce(%convert.7, %constant.6), dimensions={1}, to_apply=%region_1.2.clone.
%fused_computation.2 (param_0.9: bf16[1024,512], param_1.10: bf16[1024]) -> bf16[1024] {
%param_0.9 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0)
%convert.5 = f32[1024,512]{1,0:T(8,128)} convert(%param_0.9)
%param_1.10 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} parameter(1)
%sub.10 = bf16[1024,512]{1,0:T(8,128)(2,1)} broadcast(%param_1.10), dimensions={0}
%convert.6 = f32[1024,512]{1,0:T(8,128)} convert(%sub.10)
%sub.8 = f32[1024,512]{1,0:T(8,128)} subtract(%convert.5, %convert.6)
%exp.5 = f32[1024,512]{1,0:T(8,128)} exponential(%sub.8)
%convert.7 = bf16[1024,512]{1,0:T(8,128)(2,1)} convert(%exp.5)
%constant.6 = bf16[]{:T(256)} constant(0)
ROOT %reduce.1 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} reduce(%convert.7, %constant.6), dimensions={1}, to_apply=%region_1.2.clone
}the whole module, backend_config intact · 51 lines
HloModule jit_row_softmax, is_scheduled=true, entry_computation_layout={(bf16[1024,512]{1,0:T(8,128)(2,1)})->bf16[1024,512]{1,0:T(8,128)(2,1)}}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true}, frontend_attributes={arg_layout_modes="default",arg_memory_spaces="0",out_layout_modes="default",out_memory_spaces="0"}
%region_0.1 (reduce_max.3: bf16[], reduce_max.4: bf16[]) -> bf16[] {
%reduce_max.4 = bf16[]{:T(256)} parameter(1), metadata={op_name="reduce_max"}
%reduce_max.3 = bf16[]{:T(256)} parameter(0), metadata={op_name="reduce_max"}
ROOT %reduce_max.5 = bf16[]{:T(256)} maximum(%reduce_max.3, %reduce_max.4), metadata={op_name="jit(row_softmax)/reduce_max" stack_frame_id=20}, backend_config={"aliasing_operands":{"lists":[]},"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
}
%region_1.2.clone (reduce_sum.0: bf16[], reduce_sum.1: bf16[]) -> bf16[] {
%reduce_sum.1 = bf16[]{:T(256)} parameter(1), metadata={op_name="reduce_sum"}
%reduce_sum.0 = bf16[]{:T(256)} parameter(0), metadata={op_name="reduce_sum"}
ROOT %reduce_sum.2 = bf16[]{:T(256)} add(%reduce_sum.0, %reduce_sum.1), metadata={op_name="jit(row_softmax)/reduce_sum" stack_frame_id=20}, backend_config={"aliasing_operands":{"lists":[]},"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
}
%fused_computation.2 (param_0.9: bf16[1024,512], param_1.10: bf16[1024]) -> bf16[1024] {
%param_0.9 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0)
%convert.5 = f32[1024,512]{1,0:T(8,128)} convert(%param_0.9)
%param_1.10 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} parameter(1)
%sub.10 = bf16[1024,512]{1,0:T(8,128)(2,1)} broadcast(%param_1.10), dimensions={0}, metadata={op_name="jit(row_softmax)/sub" stack_frame_id=20}
%convert.6 = f32[1024,512]{1,0:T(8,128)} convert(%sub.10)
%sub.8 = f32[1024,512]{1,0:T(8,128)} subtract(%convert.5, %convert.6), metadata={op_name="jit(row_softmax)/sub" stack_frame_id=20}, backend_config={"flag_configs":[],"float_type_correction_info":{"original_shape":{"dimensions":["1024","512"],"element_type":"BF16","is_dynamic_dimension":[false,false],"layout":{"dim_level_types":[],"dim_ordered":[],"dim_unique":[],"dynamic_shape_metadata_prefix_bytes":"0","element_size_in_bits":"0","index_primitive_type":"PRIMITIVE_TYPE_INVALID","memory_space":"0","minor_to_major":["1","0"],"pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","split_configs":[],"tail_padding_alignment_in_elements":"1","tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}]},"tuple_shapes":[]},"original_type":"BF16"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
%exp.5 = f32[1024,512]{1,0:T(8,128)} exponential(%sub.8), metadata={op_name="jit(row_softmax)/exp" stack_frame_id=20}, backend_config={"flag_configs":[],"float_type_correction_info":{"original_shape":{"dimensions":["1024","512"],"element_type":"BF16","is_dynamic_dimension":[false,false],"layout":{"dim_level_types":[],"dim_ordered":[],"dim_unique":[],"dynamic_shape_metadata_prefix_bytes":"0","element_size_in_bits":"0","index_primitive_type":"PRIMITIVE_TYPE_INVALID","memory_space":"0","minor_to_major":["1","0"],"pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","split_configs":[],"tail_padding_alignment_in_elements":"1","tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}]},"tuple_shapes":[]},"original_type":"BF16"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
%convert.7 = bf16[1024,512]{1,0:T(8,128)(2,1)} convert(%exp.5)
%constant.6 = bf16[]{:T(256)} constant(0), metadata={op_name="jit(row_softmax)/reduce_sum" stack_frame_id=20}
ROOT %reduce.1 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} reduce(%convert.7, %constant.6), dimensions={1}, to_apply=%region_1.2.clone, metadata={op_name="jit(row_softmax)/reduce_sum" stack_frame_id=20}
}
%fused_computation.1 (param_0.8: bf16[1024,512], param_1.9: bf16[1024], param_2.4: bf16[1024]) -> bf16[1024,512] {
%param_0.8 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0)
%convert.1 = f32[1024,512]{1,0:T(8,128)} convert(%param_0.8)
%param_1.9 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} parameter(1)
%sub.9 = bf16[1024,512]{1,0:T(8,128)(2,1)} broadcast(%param_1.9), dimensions={0}, metadata={op_name="jit(row_softmax)/sub" stack_frame_id=20}
%convert.2 = f32[1024,512]{1,0:T(8,128)} convert(%sub.9)
%sub.2 = f32[1024,512]{1,0:T(8,128)} subtract(%convert.1, %convert.2), metadata={op_name="jit(row_softmax)/sub" stack_frame_id=20}, backend_config={"flag_configs":[],"float_type_correction_info":{"original_shape":{"dimensions":["1024","512"],"element_type":"BF16","is_dynamic_dimension":[false,false],"layout":{"dim_level_types":[],"dim_ordered":[],"dim_unique":[],"dynamic_shape_metadata_prefix_bytes":"0","element_size_in_bits":"0","index_primitive_type":"PRIMITIVE_TYPE_INVALID","memory_space":"0","minor_to_major":["1","0"],"pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","split_configs":[],"tail_padding_alignment_in_elements":"1","tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}]},"tuple_shapes":[]},"original_type":"BF16"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
%exp.3 = f32[1024,512]{1,0:T(8,128)} exponential(%sub.2), metadata={op_name="jit(row_softmax)/exp" stack_frame_id=20}, backend_config={"flag_configs":[],"float_type_correction_info":{"original_shape":{"dimensions":["1024","512"],"element_type":"BF16","is_dynamic_dimension":[false,false],"layout":{"dim_level_types":[],"dim_ordered":[],"dim_unique":[],"dynamic_shape_metadata_prefix_bytes":"0","element_size_in_bits":"0","index_primitive_type":"PRIMITIVE_TYPE_INVALID","memory_space":"0","minor_to_major":["1","0"],"pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","split_configs":[],"tail_padding_alignment_in_elements":"1","tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}]},"tuple_shapes":[]},"original_type":"BF16"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
%param_2.4 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} parameter(2)
%div.1 = bf16[1024,512]{1,0:T(8,128)(2,1)} broadcast(%param_2.4), dimensions={0}, metadata={op_name="jit(row_softmax)/div" stack_frame_id=20}
%convert.3 = f32[1024,512]{1,0:T(8,128)} convert(%div.1)
%div.0 = f32[1024,512]{1,0:T(8,128)} divide(%exp.3, %convert.3), metadata={op_name="jit(row_softmax)/div" stack_frame_id=20}, backend_config={"flag_configs":[],"float_type_correction_info":{"original_shape":{"dimensions":["1024","512"],"element_type":"BF16","is_dynamic_dimension":[false,false],"layout":{"dim_level_types":[],"dim_ordered":[],"dim_unique":[],"dynamic_shape_metadata_prefix_bytes":"0","element_size_in_bits":"0","index_primitive_type":"PRIMITIVE_TYPE_INVALID","memory_space":"0","minor_to_major":["1","0"],"pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","split_configs":[],"tail_padding_alignment_in_elements":"1","tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}]},"tuple_shapes":[]},"original_type":"BF16"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
ROOT %convert.4 = bf16[1024,512]{1,0:T(8,128)(2,1)} convert(%div.0)
}
ENTRY %main.3 (x.1: bf16[1024,512]) -> bf16[1024,512] {
%constant.3 = bf16[]{:T(256)} constant(-inf)
%x.1 = bf16[1024,512]{1,0:T(8,128)(2,1)} parameter(0), metadata={op_name="x"}
%copy-start = (bf16[1024,512]{1,0:T(8,128)(2,1)S(1)}, bf16[1024,512]{1,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%x.1), backend_config={"dma_priority":1,"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
%reduce_max.7 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} reduce(%x.1, %constant.3), dimensions={1}, to_apply=%region_0.1, metadata={op_name="jit(row_softmax)/reduce_max" stack_frame_id=20}, backend_config={"aliasing_operands":{"lists":[]},"flag_configs":[],"retry_config":{"retry_count":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"1052672"}],"window_config":{"buffering_level":"2","cost_model_type":"COST_MODEL_TYPE_INVALID","estimated_bundle_count":"0","estimated_cycles":"5604","estimated_scoped_vmem_bytes":"0","estimated_vmem_bytes":"0","input_window_bounds":[],"is_mask":false,"iteration_bounds":["1","1"],"kernel_window_bounds":[],"ml_estimated_microseconds":0,"output_window_bounds":["128","4"],"pad_input_on_minor_dim":"0","pad_output_on_minor_dim":"0"}}
%copy-done = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} copy-done(%copy-start), backend_config={"dma_priority":1,"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
%fusion.1 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} fusion(%copy-done, %reduce_max.7), kind=kLoop, calls=%fused_computation.2, metadata={op_name="jit(row_softmax)/reduce_sum" stack_frame_id=20}, backend_config={"aliasing_operands":{"lists":[]},"flag_configs":[],"retry_config":{"retry_count":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"1146880"}],"window_config":{"buffering_level":"2","cost_model_type":"COST_MODEL_TYPE_INVALID","estimated_bundle_count":"0","estimated_cycles":"6212","estimated_scoped_vmem_bytes":"0","estimated_vmem_bytes":"0","input_window_bounds":[],"is_mask":false,"iteration_bounds":["1","1"],"kernel_window_bounds":[],"ml_estimated_microseconds":0,"output_window_bounds":["128","4"],"pad_input_on_minor_dim":"0","pad_output_on_minor_dim":"0"}}
ROOT %fusion = bf16[1024,512]{1,0:T(8,128)(2,1)} fusion(%copy-done, %reduce_max.7, %fusion.1), kind=kLoop, calls=%fused_computation.1, metadata={op_name="jit(row_softmax)/div" stack_frame_id=20}, backend_config={"aliasing_operands":{"lists":[]},"flag_configs":[],"retry_config":{"retry_count":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"1196032"}],"window_config":{"buffering_level":"2","cost_model_type":"COST_MODEL_TYPE_INVALID","estimated_bundle_count":"0","estimated_cycles":"8740","estimated_scoped_vmem_bytes":"0","estimated_vmem_bytes":"0","input_window_bounds":[],"is_mask":false,"iteration_bounds":["1","1"],"kernel_window_bounds":[],"ml_estimated_microseconds":0,"output_window_bounds":["128","4"],"pad_input_on_minor_dim":"0","pad_output_on_minor_dim":"0"}}
}What a reduction does to a fusion
The computation above is worth one more pass, because it shows the general shape of what fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → buys on a memory-bound program. Six instructions went inside: a convert up to f32, a broadcast of the row maxima, a convert, a subtract, an exponential, a convert back down to bf16. Then a reduce. Every one of those intermediates is a full 1024 by 512 array that now never exists in memory, because the loop that computes them feeds the reduction directly.
The reduction is the root, and that is the pattern rather than a coincidence. This repo's own note on the softmax capture says it plainly: reductions become fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → roots, and the max and the sum each anchor one. A reduction is a natural stopping point because its output is small; everything upstream of it that is elementwise can be pulled in without the fused kernel having to materialize anything wide.
Notice also what did not get merged. %fused_computation.1, the ROOT fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla →, recomputes the same subtract and the same exponential that %fused_computation.2 already computed. The compiler chose to do the exponential twice rather than write a 1024 by 512 f32 intermediate to memory and read it back. That is the cost model deciding that arithmetic is cheaper than traffic, visible in the dump, on a program small enough to check by hand.
The three ways a merge does not happen
When two operations you expected to fuse did not, there are three distinct reasons, and telling them apart is most of the diagnostic skill. The first is structural. The pass can only merge instructions that already sit on a producer-consumer edge, so if two values are related in your head but not connected in the graph, no fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → pass has anything to consider. This is the wall the path spends a whole stage on at /l/xla and returns to at /xla/fusion, and it is why the naive softmax spill survives every pass in the pipeline.
The second is profitability. The candidate existed, the queue ranked it, and the difference of time estimates came out negative or too small. These are the ones the FusionProcessDumpProto can tell you about directly, and the ones the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → x-ray at /gym/xla#fusion is built out of: three real before-and-after pairs where the ranking made a call you can check against your own intuition.
The third is that something sits between the two operations that is not fusible at all. The softmax capture has a clean instance. Between the entry parameter and the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → that reads it there is a copy-start and copy-done pair whose only effect is to move the value into memory space S(1). A fusion cannot absorb a memory-space move, so the copy stays outside, and the fusion's parameter is the copy-done rather than the parameter itself. Bitcasts get the opposite treatment in the same capture: %bitcast_fusion is a fusion whose entire body is one bitcast, which exists because the backend wanted a fusion-shaped node in that position.
Check yourself
01 How do you match a fusion instruction to the computation it runs?
Only through its calls= field. The numeric suffixes are independent counters, and matching fusion.N to fused_computation.N by number will be wrong more often than right.
02 What are the three ways a merge does not happen?
No producer-consumer edge exists in the graph; the priority queue judged the merge unprofitable; or something unfusable sits between the two, like a memory-space copy.
Readings
- FusionKind in hlo_instruction.h ↗ the four kinds and the comment describing each
- priority_fusion.h ↗ the priority formula, in the header that implements it
- HLO operation semantics ↗ the op contracts fusibility checks are built on