OpenXLA’s effectiveness as a portability layer hinges on a multi-level dialect hierarchy within the MLIR framework. This post summarizes how the hierarchy is organized and why the design choices matter.

StableHLO: the portability contract

StableHLO is the bridge between frontend frameworks and backend compilers such as XLA and IREE. Any framework capable of producing StableHLO programs is compatible with any compiler that can consume them.

The specification defines approximately 100 operations and commits to long-term stability: five years of backward compatibility and two years of forward compatibility. Stability is managed via VHLO (Versioned StableHLO), an add-only dialect that snapshots the StableHLO dialect at specific points in time, so consumers can deserialize and upgrade payloads from older versions of the stack.

CHLO: client-level abstractions

The CHLO (Client HLO) dialect aligns with the API surface of XlaBuilder in C++. It models “syntactic sugar” and high-level mathematical compositions before they are materialized into lower-level dialects. This hierarchy enables hardware-independent program simplification and refines dynamically-shaped programs using concrete input arguments.

Quantization semantics

StableHLO’s types emphasize domain-specific tensor requirements. Quantized element types are governed by rigorous constraints:

• storage_type (integer), expressed_type (floating-point), scales (floating-point constants).
• For per-axis quantization, quantization_dimension < rank(T).
• Constraints C1–C13 in the spec enforce that storage_min/storage_max fit within the storage type and that scales are strictly positive and finite.

Enforcing these semantics at the compiler level prevents numerical issues common in lower-precision data formats.

Shardy: unified tensor partitioning

Shardy is an MLIR-based partitioning system from the merged GSPMD and PartIR teams. It provides a unified partitioning dialect integrated into the StableHLO layer, delivering consistent SPMD partitioning semantics across all OpenXLA backends. Shardy is directly relevant to the paper’s H3 hypothesis on SPMD scaling efficiency.

Deprecation note

The internal MHLO (Meta-HLO) dialect is being deprecated. Useful passes (canonicalization, folder patterns) are migrating into StableHLO itself, ensuring that all hardware-independent graph simplifications are available to every PJRT plugin regardless of target compiler IR.

See the full discussion in the paper, Section 2.