Two representations, one with a grammar
The path separates these two at /xla/ingestion, and the chapter above this lesson reads HLO; the split they teach is correct: StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → is the wire format, HLO is the workbench. What that pairing leaves out is how differently the two are specified, which turns out to matter every time you go looking for the rules.
StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → has a spec with a grammar in it. Programs are a sequence of functions, Program ::= {Func}, each function written func.func @name(inputs) -> (outputs) { body }. Operations are named stablehlo.mnemonic, they consume values and attributes and produce outputs, and the type system is written down: tensor types have a shape and an element type, shapes cannot be unranked, and quantized element types map an integer storage range onto an expressed floating-point type. Versioning lives in a separate dialect, VHLO, which is what lets a serialized program survive a compiler upgrade.
HLO has none of that. There is no published grammar for the text format, no versioned dialect, and no stability promise about how a module prints. What HLO has instead is a C++ class hierarchy and a set of invariants those classes hold. HloModule owns its computations and names one of them the entry. HloComputation owns an ordered list of instructions and marks exactly one ROOT. Every HloInstruction carries a Shape from the moment it is constructed. Passes are allowed to rewrite anything they like as long as those three statements stay true afterwards.
module @jit_naive_attention attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func public @main(%arg0: tensor<1024x128xbf16>, %arg1: tensor<1024x128xbf16>, %arg2: tensor<1024x128xbf16>) -> (tensor<1024x128xbf16> {jax.result_info = "result"}) {
%0 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<1024x128xbf16>) -> tensor<128x1024xbf16>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1024x128xbf16>, tensor<128x1024xbf16>) -> tensor<1024x1024xbf16>
%cst = stablehlo.constant dense<0xFF80> : tensor<bf16>
%2 = stablehlo.reduce(%1 init: %cst) applies stablehlo.maximum across dimensions = [1] : (tensor<1024x1024xbf16>, tensor<bf16>) -> tensor<1024xbf16>the whole unoptimized module · 21 lines
module @jit_naive_attention attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func public @main(%arg0: tensor<1024x128xbf16>, %arg1: tensor<1024x128xbf16>, %arg2: tensor<1024x128xbf16>) -> (tensor<1024x128xbf16> {jax.result_info = "result"}) {
%0 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<1024x128xbf16>) -> tensor<128x1024xbf16>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1024x128xbf16>, tensor<128x1024xbf16>) -> tensor<1024x1024xbf16>
%cst = stablehlo.constant dense<0xFF80> : tensor<bf16>
%2 = stablehlo.reduce(%1 init: %cst) applies stablehlo.maximum across dimensions = [1] : (tensor<1024x1024xbf16>, tensor<bf16>) -> tensor<1024xbf16>
%3 = stablehlo.broadcast_in_dim %2, dims = [0] : (tensor<1024xbf16>) -> tensor<1024x1xbf16>
%4 = stablehlo.broadcast_in_dim %3, dims = [0, 1] : (tensor<1024x1xbf16>) -> tensor<1024x1024xbf16>
%5 = stablehlo.subtract %1, %4 : tensor<1024x1024xbf16>
%6 = stablehlo.exponential %5 : tensor<1024x1024xbf16>
%7 = stablehlo.convert %6 : (tensor<1024x1024xbf16>) -> tensor<1024x1024xf32>
%cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%8 = stablehlo.reduce(%7 init: %cst_0) applies stablehlo.add across dimensions = [1] : (tensor<1024x1024xf32>, tensor<f32>) -> tensor<1024xf32>
%9 = stablehlo.broadcast_in_dim %8, dims = [0] : (tensor<1024xf32>) -> tensor<1024x1xf32>
%10 = stablehlo.convert %9 : (tensor<1024x1xf32>) -> tensor<1024x1xbf16>
%11 = stablehlo.broadcast_in_dim %10, dims = [0, 1] : (tensor<1024x1xbf16>) -> tensor<1024x1024xbf16>
%12 = stablehlo.divide %6, %11 : tensor<1024x1024xbf16>
%13 = stablehlo.dot_general %12, %arg2, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1024x1024xbf16>, tensor<1024x128xbf16>) -> tensor<1024x128xbf16>
return %13 : tensor<1024x128xbf16>
}
}What a name promises
Read a few HLO computations and the shape of the invariant set shows up in the syntax itself. A name is defined once, on the left of exactly one instruction, and referenced by name everywhere it is used. That single-definition property is what makes dataflow analysis a matter of following names rather than reasoning about mutation, and it is why the path can say at /xla/hlo 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 → depends on dataflow analysis without that sounding like a heavy claim.
The second promise is that no pass ever has to infer a type. The shape is printed on every instruction because it is stored on every instruction. A pass that wants to know whether two operands match reads two Shape objects and compares them; it does not walk backwards through the graph reconstructing what something must be. That is a real design decision with a cost, since every rewrite has to compute and set the new shape explicitly, and a real payoff, since a partially rewritten module is still fully typed.
The third is that structure nests through named computations rather than through blocks. A reduce carries to_apply=%region_0.1 and the region is a separate computation printed alongside. 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 → carries calls=%fused_computation.1 the same way. So a computation you find in a dump is either the entry, or something named by an instruction elsewhere, and searching the module for its name finds the caller in one step.
%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}
}the whole optimized softmax module · 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"}}
}Decoding a shape all the way down
The path teaches the layout brace at /xla/layout-memory and stops at the minor-to-major list, which is the part that changes strides. A real TPU dump prints more inside that brace, and the rest of it is readable too. Take bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} and pull it apart left to right.
The interesting fragments are the two after the colon. The tiles are a physical blocking of the last dimensions, and the second one is easy to identify from the capture itself rather than from documentation: every bf16 shape in this module carries T(8,128)(2,1) and every f32 shape carries T(8,128) alone. One tile appears only on the 16-bit type, which is what a packing tile for two 16-bit values looks like. Scalars get a different tile entirely, T(256), as the reduce region above shows.
S(1) is a memory space, and the capture proves it is not decoration. In the softmax module a copy-start takes %x.1 in the default space and produces a shape that differs from its input in exactly one way, the trailing S(1); a matching copy-done completes it. Same element type, same dimensions, same layout, same tiles. A copy whose only visible effect is the memory space is a move between memories, and it is scheduled as an async pair because it takes real time.
| fragment | what it says |
|---|---|
| bf16 | element type |
| [1024,512] | logical dimensions, major to minor as written |
| {1,0 | minor-to-major order: dimension 1 varies fastest, so rows are contiguous |
| T(8,128) | physical tiling of the last two dimensions |
| (2,1) | a second tile present on every bf16 shape in this capture and on no f32 shape in it |
| S(1) | memory space 1, reached in this module through a copy-start and copy-done pair |
The header line is a contract with PJRT
Every dumped HLO module opens with one long line that most readers skip, and two of its fields are worth not skipping. entry_computation_layout={(bf16[1024,512]{1,0:T(8,128)(2,1)})->bf16[1024,512]{1,0:T(8,128)(2,1)}} is the promise the compiled program makes to whoever calls it, stated in exactly the vocabulary the previous section decoded. That promise is what a PjRtBuffer has to satisfy before Execute will accept it, which is where the pjrt lesson and this one meet.
is_scheduled=true changes how you read everything below it. A module dumped before scheduling lists instructions in an order that only has to respect dataflow; a scheduled module lists them in the order they will actually run. So in a scheduled dump the distance between a copy-start and its copy-done is not a formatting accident, it is the window the scheduler chose to overlap that transfer against, and counting what sits between them tells you how much compute is covering the move.
The rest of the header line is sharding and frontend bookkeeping: allow_spmd_sharding_propagation_to_parameters, arg_layout_modes, arg_memory_spaces. On a single-device capture these read as trivia. They stop being trivia the moment the spmd unit's lesson gets to a partitioned module, because they are the flags that decided whether the partitioner was allowed to change your inputs' shardings at all.
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"}the module under that header · 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"}}
}Fluency, drilled on one capture
Reading fluency here is not a matter of understanding, it is a matter of having done it enough times that the eye stops stalling. The three captures in this repo are enough material: naive attention, an MLP, and a row softmax, each dumped before and after optimization on a TPU v6 lite. The routine that pays off is to start at ROOT in the ENTRY computation and walk backwards through operands, which reconstructs the program in the order the compiler thinks about it.
Doing that on the softmax capture turns up a naming trap worth meeting once, deliberately, instead of by accident. The ROOT instruction is %fusion and it calls %fused_computation.1. The instruction feeding it is %fusion.1 and that one calls %fused_computation.2. Instruction suffixes and computation suffixes are independent uniquifying counters, and 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 → numbered 1 has no particular relationship to a computation numbered 1.
The instruction number and the computation number are two different counters.
The rest of what you need is repetition against real modules, which is what the drills on /gym/xla exist for. The one thing worth carrying out of this lesson in advance is the habit of decoding the full brace rather than the first two numbers, because on a TPU capture the tiles and the memory space are where the interesting decisions ended up.
Check yourself
01 What does StableHLO have that HLO does not?
A written grammar, a spec, and versioned serialization. HLO has a C++ class hierarchy and a set of invariants passes must leave standing.
02 In bf16[1024,512]{1,0:T(8,128)(2,1)S(1)}, what are the last two fragments?
The (2,1) is the packing tile that appears only on 16-bit types in the capture; S(1) is a memory space, reached through a copy-start and copy-done pair.
Readings
- the StableHLO spec ↗ the grammar and type system HLO does not have
- HLO operation semantics ↗ op by op, including dot_general dimension numbers
- XLA architecture ↗ where the two representations sit in the flow