StableHLO Documentation

StableHLO is an operation set for high-level operations (HLO) in machine learning (ML) models. It serves as a portability layer between ML frameworks (such as TensorFlow, JAX, and PyTorch) and ML compilers (such as XLA and IREE).

Key features include:

• Interoperability: Frameworks producing StableHLO programs are compatible with compilers consuming them.
• Based on MHLO: It enhances the MHLO dialect with additional functionality like serialization and versioning.
• Serialization: Uses MLIR bytecode as the serialization format.
• Compatibility: Provides backward and forward compatibility guarantees to ensure stability as the project evolves.

VHLO (Versioned StableHLO) is an add-only dialect used for serialization and stability. Unlike the standard StableHLO dialect, which only contains the latest version of operations, VHLO provides a snapshot of the StableHLO dialect at specific points in time by versioning individual program elements (ops, types, and attributes).

Key characteristics:

• Immutability: Once a feature is added, it cannot be modified in a way that impacts semantics. Any change requires a new versioned entry.
• Versioning: Each op, type, or attribute is associated with a version range (e.g., from version 0.9.0 to 0.10.0).
• Evolution Tracking: VHLO captures the evolution of an operation. For example, if an attribute is added to an operation in a newer StableHLO release, VHLO will contain both the old version (without the attribute) and the new version (with the attribute).

Several major machine learning frameworks leverage OpenXLA and StableHLO for high-performance model execution:

• JAX: A framework with a NumPy-like API for high-performance ML models.
• PyTorch/XLA: Provides a bridge from PyTorch to OpenXLA and StableHLO.
• TensorFlow: A large-scale ML framework with an extensive ecosystem.
• Reactant.jl: A Julia framework for optimizing and executing code via OpenXLA, StableHLO, and MLIR.
• GoMLX: An ML framework for the Go language, which includes gopjrt (a raw XlaBuilder+PJRT wrapper for Go tested on CPU, GPU, and TPU).

StableHLO programs follow a hierarchical structure:

1. Programs: The top-level container.
2. Functions: Programs consist of one or more StableHLO functions.
3. Ops: Functions consist of individual StableHLO operations (ops).

The specification defines the semantics of individual ops, how they execute together within a program, and the notation used to describe them.

StableHLO supports dynamically-shaped tensors. However, shapes must agree at runtime. If shapes do not match during execution, the behavior is undefined. It is the responsibility of the producer to generate code that ensures correct shapes at runtime, as StableHLO does not provide an explicit runtime shape assertion operation.

func.func @foo(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> tensor<?xi32> {
    %0 = stablehlo.add %arg0, %arg1 : tensor<?xi32>
    return %0 : tensor<?xi32>
}

The following ops are subject to change and may be moved to CHLO (Custom HLO) or removed entirely. Currently, only a 6-month compatibility guarantee is provided for these ops:

• MapOp: Likely to be removed; stablehlo.composite is the preferred way to handle region-based logic.
• EinsumOp: Likely to be moved to CHLO.
• TorchIndexSelectOp: Likely to be moved to CHLO; can be decomposed using gather.
• CrossReplicaSumOp: Being moved to CHLO (it is a sugar for all-reduce).

If a target environment does not support stablehlo.tan directly, it can be represented via decomposition. A common mathematical decomposition is:

TanOp(x) = DivOp(SineOp(x), CosineOp(x))

Note that this may introduce numerical differences compared to the native TanOp. The proposal suggests that pattern matching SineOp/CosineOp combinations into TanOp should be an opt-in pass to allow users to maintain separate operations if desired.

When reading StableHLO specifications, it is important to distinguish how names are treated in different contexts:

• Value: An actual value with a known type (e.g., a specific tensor dense<[[1, 2]]>).
• Placeholder: A future value where only the type is known, but the actual data is not yet available (typical for non-constant inputs during compile-time constraint checking).

StableHLO supports four strategies for splitting the process grid into process groups for collective operations:

1. cross_replica: Computes the Cartesian product of replica_groups and partition_ids. Each group contains processes with the same partition_id but different replica_ids within a group.
2. cross_partition: Computes the Cartesian product of partition_groups and replica_ids. Each group contains processes with the same replica_id but different partition_ids within a group.
3. cross_replica_and_partition: Computes the Cartesian product of replica_groups and partition_ids, where each group contains both replica and partition variations.
4. flattened_ids: Uses a list of lists of 'flattened' process IDs (calculated as replica_id * num_partitions + partition_id) to define groups.