Legalizes from CHLO ops flow to StableHLO and Shape ops


Legalize shape-related ops to StableHLO.

An experimental pass that legalizes shape-related ops to StableHLO ops.

Bringing shape and data computations together via an optional pass will make it possible for the StableHLO ecosystem to potentially leverage the compilation pipelines that use StableHLO operations to model dynamism.


Folds StableHLO operations


Canonicalizes StableHLO operations


Canonicalizes dynamic StableHLO ops into static ops.

Replaces dynamic StableHLO ops like DynamicReshapeOp with the corresponding static counterparts like ReshapeOp if all the dynamic elements of these ops are actually constant.

For example, if the output_shape operand of DynamicReshapeOp is a constant value, then the operation can be transformed to ReshapeOp.


Inserts probe instrumentation instructions in a StableHLO program.

Walks through a StableHLO program and inserts a probe instrumentation operation after each suitable operation (see below for how a suitable operation is defined). Instrumentation is used to extract intermediate tensor values from the StableHLO reference interpreter for later comparison with other runtimes.

All operations are considered suitable for instrumentation, except constant ops, ops which do not have any tensor return values (i.e. an op that produces a tuple or a token or no return values will not be instrumented). Suitable operations will be instrumented regardless of their level of nesting. That is, operations inside loop/branch regions will also be instrumented.

Instrumented operations will have their return values written to disk using the NumPy data format as they are executed. If the useDebugInfo pass option is enabled, location debug information will be used when available to uniquely identify instrumented tensor values (i.e. the pass will extract probe_id from NamedLoc(probe_id@<...>) and use the format probe_id.#). Otherwise, instrumented values will be referred to in the increasing sequence: probe1, probe2, ... See interpreter.probe for additional information on how data is serialized.


-useDebugInfo : Whether or not to use location debug data as `probe_id` values.


Replaces composite ops with a call to their decomposition

Replaces composite ops with a call to their decomposition, e.g. the below:

stablehlo.composite "my_namespace.my_op" %arg0, %arg1 {
  decomposition = @bar,
  version = 1,
  composite_attributes = {
    "my_attribute": "my_value"

Will become:

func.call @bar(%arg0, %arg1)

A subset of composites can be excepted from this transformation using the "except" flag, e.g.:

stablehlo-opt --stablehlo-legalize-composite-to-call=except='foo.baz,foo.qux'


-except : Names of composites that should not be replaced with calls.


Legalize deprecated ops to well-supported ops.

The StableHLO v1.0 Opset Deprecations RFC (#2283) proposes to remove several redundant ops. This pass helps to evaluate the impact of these op removals in various compilation pipelines by legalizing them to their long-term supported counterparts.


-fail-on-unused : Fail on (mostly) unused ops that are deprecated without any fallback.


Legalize StableHLO to VHLO.


Refines the argument shapes of the main function.

Modifies the arguments of the main function using the input type signature. Wraps arguments in custom_call @stablehlo.shape_refinement_operand_wrapper to keep the IR valid before shape refinement is run.

The refinedTypesOption can be used to specify a list of refined types. This can be specified in MLIR with --types='tensor<...>,tensor<...>', or passed to the pass create method. The refinement type list must specify the type of every argument to the main method being refined.


-types : The new types to be used for the main function's arguments, specified as an MLIR TypeRange 'tensor<1x2xf32>, ...'


Refines shapes across a StableHLO program.

Walks through a StableHLO program refining shapes within ops.

The flagship use case for this pass is specializing dynamically-shaped programs to static shapes. If a dynamically-shaped StableHLO program has the right structure, then updating its argument types from dynamic shapes to static shapes and running this pass will propagate static shapes across the program.


Legalize VHLO to StableHLO.


Convert between versions of VHLO.


-target : The target version. Must be a version of the form #.#.# or 'current'.