How a StableHLO Module Becomes HLO: mlir_to_hlo.cc, Line by Line

Your framework does not hand the compiler a graph. It hands it a versioned document written in a dialect called StableHLO, and 434 lines of C++ decide what that document is allowed to say. We read all of them. The includes, which split into the two vocabularies this file sits between. The thirteen MLIR passes a module walks through before anything called HLO exists: seven from a Shardy pipeline it pulls in unconditionally, six added here by name, ending with the one that copies captured constants into control-flow regions because XLA has no implicit capture. The Shardy-loses-to-GSPMD fallback at the top of the function, and the block sixty lines later that silently reads the flag it just switched off. The parser, and the one error message that names three different version problems because it genuinely cannot tell them apart. Then the write side, which is the longer half: the version clamp that keeps the lower of what a plugin asked for and what this build can produce, the allow-list that refuses to serialize any operation nobody has promised to read, and the twelve-week compatibility window written as a requirement rather than a number. The signature exhibit runs the real branches: pick a direction, flip the flags, and every pass your module walks through appears in order, at the line that adds it.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of How a StableHLO Module Becomes HLO: mlir_to_hlo.cc, Line by Line — built to be understood, not skimmed.

Questions

What is the difference between StableHLO and HLO?
HLO is XLA’s own intermediate representation, a protobuf that the optimizer, the layout assignment and the buffer assignment all work on. It makes no compatibility promise: a proto produced by one build is only reliably consumed by that build. StableHLO is an MLIR dialect with the same operation set in spirit, plus two things HLO does not have. It has a version number, and it has a serializer that can write a module at an older version on request. That is what makes a plugin possible, because a framework built this morning has to hand a program to a vendor library built three months ago. mlir_to_hlo.cc is the crossing between the two: MlirToXlaComputation runs a pass pipeline over the StableHLO module and then calls ConvertStablehloToHloWithOptions, which returns an HloModule that becomes the proto the rest of the compiler consumes.
What passes run before StableHLO is converted to HLO?
Thirteen, in one pass manager. Seven come from addSdyRoundTripExportPipeline, which mlir_to_hlo.cc adds unconditionally at line 131: a canonicalizer, mesh lifting, mesh dedup, the Shardy op export, the shard-map export, the attribute export, and the StableHLO sharding export. On a module with no Shardy in it they match nothing. Six more are added by name in the file itself: chlo-recompose-ops rebuilds CHLO that was serialized as custom calls, symbol-dce removes unreferenced functions, chlo-legalize-to-high-level-mhlo claims the CHLO operations XLA models directly such as TopK and Erf, chlo-legalize-to-stablehlo decomposes the rest, the complex math expander rewrites operations like log_plus_one into plain arithmetic, and sink-constants-to-control-flow copies implicitly captured constants into the regions that use them. That last one is not an optimization: XLA control flow is functional, so a loop body cannot refer to a value defined outside it. Only after all thirteen have run does line 175 do the actual conversion.
What does the twelve-week StableHLO compatibility window actually mean?
It means a framework can talk to a plugin up to roughly twelve weeks older than itself without either side being rebuilt, chosen because PJRT plugin vendors ship on about a quarterly cycle. GetDefaultStablehloVersion in mlir_to_hlo.cc does not hard-code a number; it asks StableHLO for the version satisfying CompatibilityRequirement::WEEK_12, so the constant moves when the dependency moves. At the commit this page reads, that resolves to 1.15.0 while StableHLO’s current version is 1.19.0. In practice the target is not the default at all when the plugin reports one: PJRT reads an attribute named stablehlo_current_version from the plugin and passes that as the requested target, falling back to the twelve-week default only when the plugin does not report a version. Shardy has an identical function with the same twelve-week rule for its own versioning scheme.
What happens if a PJRT plugin reports a newer StableHLO version than the framework?
Nothing fails, and the plugin does not get the version it asked for. SerializeUsingVersionedStablehlo never uses the requested target directly: it passes the request and the framework’s own current version to getSmallerVersion and writes the artifact at whichever is lower. So a plugin ahead of the framework receives the newest artifact the framework can produce, which it can read, because a newer StableHLO build can always read an older artifact. The source comment names the exact situation it defends against: usually the plugin is older, but occasionally a plugin nightly build is compiled against the latest public release of a framework. The upstream test asserts the behaviour by requesting version 100.99.98 and expecting an artifact stamped with the framework version. The only failure available here is a version string that does not parse, which returns “Invalid StableHLO target version requested.”
Why does serialization fail with "found unstable op"?
Because the module contains an operation from a dialect nobody has promised can be read back by an older build. When mixed serialization is allowed, which is what the public Serialize entry point always requests, FindPotentiallyUnstableDialects walks every operation in the module and stops at the first one that is not from StableHLO, CHLO or Shardy, is not one of exactly four func operations (the module, a function, a call and a return), and is not the single hard-coded dialect name mpmd. The error names the offending operation rather than its dialect, because the operation name is what you can act on. The upstream tests show the shapes that catch people out: func.constant together with call_indirect are core MLIR yet are outside the four-operation allow-list, and an unregistered operation fails the same way. A portable artifact is a promise that an older build can read these bytes, and that promise can only be made about dialects whose versioning is understood.

Related explainers