The HLO family tree

Lower a single arctangent and JAX hands back a module carrying three different prefixes. func.func is plain upstream MLIR. chlo.atan is an op the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → spec does not contain at all. And bolted onto the module itself are mhlo.num_partitions and mhlo.num_replicas, named for a dialect whose own README says it is deprecated and slated for removal.

None of that is cruft to route around. The five names in this family, HLO and MHLO and StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → and VHLO and CHLO, are not five levels of one stack; they are five answers to five different problems, and they arrived in the order the problems did. Read them as history and the prefixes on this dump stop being noise.

module @jit_atan_fn attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<4xf32>) -> (tensor<4xf32> {jax.result_info = ""}) {
    %0 = chlo.atan %arg0 : tensor<4xf32> -> tensor<4xf32>
    return %0 : tensor<4xf32>
  }
}

module @jit_atan_fn attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<4xf32>) -> (tensor<4xf32> {jax.result_info = ""}) {
    %0 = chlo.atan %arg0 : tensor<4xf32> -> tensor<4xf32>
    return %0 : tensor<4xf32>
  }
}

HloModule jit_atan_fn, is_scheduled=true, entry_computation_layout={(f32[4]{0})->f32[4]{0}}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true}

%fused_computation (param_0: f32[4]) -> f32[4] {
  %param_0 = f32[4]{0} parameter(0)
  %constant.0 = f32[] constant(1)
  %broadcast.0 = f32[4]{0} broadcast(f32[] %constant.0), dimensions={}
  ROOT %atan2.0 = f32[4]{0} atan2(f32[4]{0} %param_0, f32[4]{0} %broadcast.0), metadata={op_name="jit(atan_fn)/jit(main)/atan" source_file="<string>" source_line=4}
}

ENTRY %main.5 (Arg_0.1: f32[4]) -> f32[4] {
  %Arg_0.1 = f32[4]{0} parameter(0), metadata={op_name="v"}
  ROOT %broadcast_atan2_fusion = f32[4]{0} fusion(f32[4]{0} %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(atan_fn)/jit(main)/atan" source_file="<string>" source_line=4}
}


module @jit_erf_fn attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<4xf32>) -> (tensor<4xf32> {jax.result_info = ""}) {
    %0 = chlo.erf %arg0 : tensor<4xf32> -> tensor<4xf32>
    return %0 : tensor<4xf32>
  }
}

HloModule jit_erf_fn, is_scheduled=true, entry_computation_layout={(f32[4]{0})->f32[4]{0}}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true}

ENTRY %main.3 (Arg_0.1: f32[4]) -> f32[4] {
  %Arg_0.1 = f32[4]{0} parameter(0), metadata={op_name="v"}
  ROOT %erf.2 = f32[4]{0} erf(f32[4]{0} %Arg_0.1), metadata={op_name="jit(erf_fn)/jit(main)/erf" source_file="<string>" source_line=5}
}

HLO is the oldest thing in this picture by years, and it was never a dialect of anything. It is an enum. xla/hlo/ir/hlo_opcode.h lists 132 opcodes as macro entries of the form V(kAtan2, "atan2", 2), name and printed spelling and arity, and the comment above the list calls them "High-level optimizer instruction opcodes" and describes them as "a flattened form of the UnaryOp, BinaryOp, ... opcodes present in the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → service protobuf".

That origin explains what HLO is missing. It was the internal vocabulary of one compiler talking to itself, so it needed no written grammar, no version number, and no story for what happens when the compiler upgrades under a saved file. The xla:hlo unit takes apart the invariants HLO offers instead of a spec; what matters here is only the date. Everything else in this family was built after HLO, around HLO, or to replace a job HLO was doing badly.

MLIR gave compiler authors a way to keep several op vocabularies inside one module and lower between them under shared infrastructure. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → wanted its ops in that world, so the ops got an MLIR dialect: MHLO, still sitting in the tree at xla/mlir_hlo/mhlo, with 138 op definitions that mirror the HLO opcodes. For a few years MHLO was the bridge every frontend crossed to reach XLA.

That bridge is closed now. The README on that directory opens with a deprecation notice and a line that settles the question of whether you should learn MHLO in any depth.

Users of MHLO should migrate to StableHLO whenever possible.

The dialect is going away and the prefix is not, which is why your lowered module still says mhlo.num_partitions. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → keeps a single header listing every attribute a StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → module is allowed to carry across the border, and the mhlo. namespace is where the facts live that the StableHLO spec has no field for: replica counts, SPMD shardings, entry layouts, buffer donation. An attribute outside that list does not survive the trip.

namespace xla {

// This file captures all discardable attributes that XLA supports.
// Attributes not in this list will be dropped when exporting to StableHLO.

// Module level attributes require namespacing.
inline constexpr char kMhloCrossProgramPrefetches[] =
    "mhlo.cross_program_prefetches";
inline constexpr char kMhloInputOutputAlias[] = "mhlo.input_output_alias";
inline constexpr char kMhloSpmdOutputSharding[] = "mhlo.spmd_output_sharding";
inline constexpr char kMhloNumPartitions[] = "mhlo.num_partitions";
inline constexpr char kMhloNumReplicas[] = "mhlo.num_replicas";

The openxla/stablehlo repository opens with an empty initial commit on 3 August 2022. Two weeks later the second commit lands, and its title is the whole design decision: "Bootstrap StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → from CHLO/MHLO (#1)". StableHLO did not invent an opset. It copied the one MHLO already had, from a compiler that had been running it in production for years, and then made a promise about it.

The commit message is worth reading because it names the problem in the language of the moment it was written. Frameworks and compilers were passing dialects between repositories with no agreement about what a version bump was allowed to break, and the proposal on the table was a shallow dialect that producers could vendor and upgrade under stated compatibility windows.

Bootstrap StableHLO from CHLO/MHLO (#1)

Recent discussions highlight an acute need for stability of interchange
dialects in between ML frameworks and ML compilers in the opensource
community.

In a Discourse post, a potential solution was called out: something called
"shallow dialects" that producers could vendor into their repositories and
upgrade with well-defined backward compatibility windows.

I think this presents a great opportunity for StableHLO: to start as shallow
MHLO which will bootstrap us from a well-understood baseline and will enable
us to provide a service to the community right away - backward compatibility
guarantees for MHLO (and its sister dialect CHLO as well).

A promise about serialized bytes needs somewhere to keep the old shapes of things, because the live dialect only ever holds the current shape. That somewhere is VHLO, added to the repo in December 2022, and its dialect definition describes itself as a shallow versioned copy of StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → simplified down to a bare minimum, used for upgrades, downgrades, and serialization. Every op in it carries a version range instead of a definition that can be edited.

Watch one change land. StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → 1.5.0 made the collectives variadic so all_gather could take several operands at once. The op in StableHLO was rewritten in place; the op in VHLO could not be, so all_gather_v1 had its upper bound closed at 1.4.0 and a new all_gather_v2 opened at 1.5.0. Both definitions are still in the file, side by side, differing by one Variadic<> wrapper.

Count the file and the growth shows: 142 op definitions in VHLO against 117 in StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →, and 23 of the VHLO ones have a closed upper bound. Those 23 are ops nothing will ever emit again. They stay because a reader opening a three-year-old artifact needs something to deserialize it into, which is the mechanism behind the compatibility windows the previous lesson quoted.

// VhloDialect.td, version log
//   1.5.0: Make collective ops (`all_reduce`, `all_gather`, `all_to_all`) variadic.

def VHLO_AllGatherOpV1 : VHLO_Op<"all_gather_v1", "0.9.0", "1.4.0"> {
  let arguments = (ins
    VHLO_AnyType:$operand,
    VHLO_AnyAttr:$all_gather_dim,
    VHLO_AnyAttr:$replica_groups,
    VHLO_AnyAttr:$channel_id,
    VHLO_AnyAttr:$use_global_device_ids
  );
  let results = (outs VHLO_AnyType:$result);
}

def VHLO_AllGatherOpV2 : VHLO_Op<"all_gather_v2", "1.5.0", "current"> {
  let arguments = (ins
    Variadic<VHLO_AnyType>:$operands,
    VHLO_AnyAttr:$all_gather_dim,
    VHLO_AnyAttr:$replica_groups,
    VHLO_AnyAttr:$channel_id,
    VHLO_AnyAttr:$use_global_device_ids
  );
  let results = (outs Variadic<VHLO_AnyType>:$results);
}

You can watch VHLO do its job from Python without leaving this machine. jaxlib.mlir.dialects.stablehlo exposes StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →'s own serializer, and it takes a target version: the version of the reader you are writing for, not the version you are running. Ask it for a 0.9.0 reader and the same four-op module comes out at 575 bytes; ask for 1.0.0 and it comes out at 459. Both artifacts name their target in a producer string right after the MLIR bytecode magic.

Ask for a version this build does not have and it refuses rather than guessing. The message comes from VhloToVersion.cpp, which is also where the harder failure lives: when a program uses an op that has no definition at the requested version, serialization fails with failed to convert VHLO to v<version> instead of emitting bytes the old reader would misread. Refusing to downgrade is the promise working, not the promise breaking.

print('current', shlo.get_current_version(), 'minimum', shlo.get_minimum_version())
for target in ('0.9.0', '1.0.0', '1.8.7'):
    blob = shlo.serialize_portable_artifact_str(module_text, target)
    print(target, len(blob), blob[:4], blob[5:21])
shlo.serialize_portable_artifact_str(module_text, '1.9.0')

current 1.8.7 minimum 0.9.0
0.9.0 575 b'ML\xefR' b'StableHLO_v0.9.0'
1.0.0 459 b'ML\xefR' b'StableHLO_v1.0.0'
1.8.7 459 b'ML\xefR' b'StableHLO_v1.8.7'
loc("-":1:1): error: target version 1.9.0 is greater than current version 1.8.7
ValueError: failed to serialize module

That leaves the op the lesson opened on. CHLO is older than StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → and came along in the same bootstrap; its dialect description says it models the API surface of the XlaBuilder C++ API, and that whenever the client library uses syntactic sugar or a composition of several ops for one API call, CHLO models the call and supplies conversion patterns to materialize it into lower level dialects. StableHLO has atan2 and no atan, so jnp.arctan has nowhere to land except CHLO.

The conversion pattern for it is three lines of TableGen and the comment above it states the identity outright: atan(x) = atan2(x, 1). Compile that same function on CPU and the prediction is sitting in the HLO: a constant(1) broadcast to the operand shape, then atan2(f32[4]{0} %param_0, f32[4]{0} %broadcast.0), and the frontend name survives in the instruction metadata as op_name="jit(atan_fn)/jit(main)/atan".

Now do the same with jax.scipy.special.erf and something different happens. It lowers to chlo.erf, and the compiled entry computation reads the parameter straight into one instruction, ROOT %erf.2 = f32[4]{0} erf(f32[4]{0} %Arg_0.1). No constant, 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 →, no decomposition at all. The reason is that HLO has an erf opcode and no atan opcode, and XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → ingestion knows it: the function StablehloToMhlo in xla/hlo/translate/stablehlo.cc runs a CHLO to high-level MHLO pass first, summarized in its own definition as legalizing "CHLO's with XLA counterparts, like TopK and Erf", and only then decomposes whatever is left into spec ops.

Which fate an op meets is decided by whether HLO already had an opcode for it.

So the family sorts by the problem each member was built for, and the whole tree is legible from that one dump. Here they are on one card.

// Express `atan` as
//   atan(x) = atan2(x, 1)
def : Pat<(CHLO_AtanOp NonComplexElementType:$input),
  (StableHLO_Atan2Op
    $input,
    (StableHLO_ConstantLike<"1"> $input)
  )>;

def : Pat<(CHLO_TanOp $input),
          (StableHLO_TanOp $input, ConstDefaultResultAccuracyAttr)>;