OpenXLA Project

`-chlo-legalize-to-stablehlo`

将 CHLO 操作流程合法化为 StableHLO 和 Shape 操作

`-shape-legalize-to-stablehlo`

将与形状相关的操作合法化为 StableHLO。

一种实验性传递，可将与形状相关的操作合法化为 StableHLO 操作。

通过可选的传递将形状和数据计算结合在一起，StableHLO 生态系统就有可能利用使用 StableHLO 操作来模拟动态性的编译流水线。

`-stablehlo-canonicalize-dynamism`

将动态 StableHLO 操作规范化为静态操作。

如果动态 StableHLO 操作（例如 DynamicReshapeOp）的所有动态元素实际上都是常量，则将这些操作替换为相应的静态对等项，例如将 DynamicReshapeOp 替换为 ReshapeOp 或将 DynamicBroadcastInDim 替换为 BroadcastInDim。

  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>

  ==>

  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>

`-stablehlo-check-shape-assertions`

_检查 stablehlo.custom_call @shape 断言操作。

验证 shape_assertion 自定义调用。

形状断言用于验证 StableHLO 中动态维度的限制。例如，如果框架需要强制执行 DimA < 2 的限制，则可以发出以下 IR：

%dimA = <get_dimension_size or input arg> : tensor<i32>
%c2 = stablehlo.constant dense<2> : tensor<i32>
%is_lt = stablehlo.compare LT %dimA, %c2 : tensor<i1>
stablehlo.custom_call @shape_assertion(%is_lt) { error_message = "DimA must be less than 2" }

通过测试后，如果形状正确，系统会移除 stablehlo.custom_call。

选项

-enable-shape-assertions : Whether shape assertions may generate errors.

`-stablehlo-compatibility-expander`

StableHLO 操作的兼容性扩展器。

StableHLO 操作在最新版本中会获得更新或引入新的操作。此选择启用传递通过将较新的 StableHLO 操作分解为旧版本支持的等效操作，从而扩展了与旧版 StableHLO 的向后兼容性。

为什么此卡是选择性购买的？

有时，StableHLO 操作增强功能用于极大地简化 OpenXLA 生态系统中某些常见模式的处理。这包括 TanOp（具有较高的框架和编译器支持）以及可以使用切片表示但会使分片变得更加困难的收集/分散批处理维度。对于此类新功能，我们不提供自动降级，因为这可能会舍弃后续优化中使用的重要信息。此传递可用于根据目标版本扩展这些操作，以最大限度地提高兼容性，但可能会牺牲编译的优化程度。

func.func @tan_op_non_complex(%arg0: tensor<4xf64>) -> tensor<4xf64> {
  %1 = stablehlo.tan %arg0 : tensor<4xf64>
  func.return %1 : tensor<4xf64>
}

==>

func.func @tan_op_non_complex(%arg0: tensor<4xf64>) -> tensor<4xf64> {
  %0 = stablehlo.sine %arg0 : tensor<4xf64>
  %1 = stablehlo.cosine %arg0 : tensor<4xf64>
  %2 = stablehlo.divide %0, %1 : tensor<4xf64>
  return %2 : tensor<4xf64>
}

选项

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

`-stablehlo-complex-math-expander`

StableHLO 复数数学运算的扩展器。

StableHLO 复数数学运算是使用 StableHLO 实数数学运算进行的分解。

此声明基于以下假设：不存在原生支持复数或复杂数学运算的硬件。这意味着，编译器可能实现的复杂数学运算回退机制是多余的。启用此 pass 后，所有 StableHLO 复数数学运算都将展开。

func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
  %1 = stablehlo.sqrt %arg0 : tensor<4xcomplex<f64>>
  func.return %1 : tensor<4xcomplex<f64>>
}

==>

func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
  TBD
  return %2 : tensor<4xcomplex<f64>>
}

`-stablehlo-convert-to-signless`

传递以转换 IR，使其采用无符号整数。

`-stablehlo-legalize-composite-to-call`

将复合操作替换为对其分解的调用。

将复合操作替换为对其分解的调用，例如：

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

将变为：

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

可以使用“except”标志将部分复合项排除在此转换之外，例如：

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

选项

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

`-stablehlo-legalize-deprecated-ops`

将已弃用的操作合法化为受支持的操作。

StableHLO v1.0 Opset 弃用 RFC (#2283) 提议移除几个冗余操作。此传递通过将这些操作合法化为长期支持的对应项，有助于评估这些操作移除在各种编译流水线中的影响。

选项

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

`-stablehlo-legalize-qdq-to-quantized-op`

将融合（反量化、浮点运算和量化）模式纳入 StableHLO 量化运算

将（反量化、浮点运算和量化）模式融合到 StableHLO 量化运算中注意：此 pass 不会删除任何预先存在的运算。例如，以下程序

func.func @add(%arg0: tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<16x16xf32>
  %1 = stablehlo.abs %0 : tensor<16x16xf32>
  %2 = stablehlo.uniform_quantize %1 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
  func.return %2 : tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
}

将变为：

func.func @add(%arg0: tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>) -> tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>) -> tensor<16x16xf32>
  %1 = stablehlo.abs %0 : tensor<16x16xf32>
  %2 = stablehlo.abs %arg0 : tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
  %3 = stablehlo.uniform_quantize %1 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
  return %2 : tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
}

`-stablehlo-legalize-quant-to-math`

将 StableHLO 量化操作转换为 StableHLO 原始数学运算。

将使用 UniformQuantized 类型的 StableHLO 程序转换为语义上等效的整数数学运算。

func.func @add(%arg0: tensor<!quant.uniform<i8:f32,1.0:0>>, %arg1: tensor<!quant.uniform<i8:f32,2.0:1>>) ->  tensor<!quant.uniform<i8:f32,3.0:2>> {
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32,1.0:0>>, tensor<!quant.uniform<i8:f32,2.0:1>>) -> tensor<!quant.uniform<i8:f32,3.0:2>>
  func.return %0 : tensor<!quant.uniform<i8:f32,3.0:2>>
}

将变为：

func.func @add(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
  %0 = stablehlo.convert %arg0 : (tensor<i8>) -> tensor<f32>
  %cst = stablehlo.constant dense<0.333333343> : tensor<f32>
  %1 = chlo.broadcast_multiply %0, %cst : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<f32>
  %2 = chlo.broadcast_add %1, %cst_0 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %3 = stablehlo.round_nearest_even %2 : tensor<f32>
  %4 = stablehlo.convert %3 : (tensor<f32>) -> tensor<i32>
  %5 = stablehlo.convert %arg1 : (tensor<i8>) -> tensor<f32>
  %cst_1 = stablehlo.constant dense<0.666666686> : tensor<f32>
  %6 = chlo.broadcast_multiply %5, %cst_1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_2 = stablehlo.constant dense<1.33333337> : tensor<f32>
  %7 = chlo.broadcast_add %6, %cst_2 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %8 = stablehlo.round_nearest_even %7 : tensor<f32>
  %9 = stablehlo.convert %8 : (tensor<f32>) -> tensor<i32>
  %c = stablehlo.constant dense<2> : tensor<i32>
  %10 = chlo.broadcast_add %4, %9 : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %11 = chlo.broadcast_subtract %10, %c : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %c_3 = stablehlo.constant dense<-128> : tensor<i32>
  %c_4 = stablehlo.constant dense<127> : tensor<i32>
  %12 = stablehlo.clamp %c_3, %11, %c_4 : tensor<i32>
  %13 = stablehlo.convert %12 : (tensor<i32>) -> tensor<i8>
  return %13 : tensor<i8>
}

`-stablehlo-legalize-quantized-op-to-qdq`

将量化的 StableHLO 操作分解为（反量化、浮点运算和量化）模式。

使用统一的量化/反量化操作分解 StableHLO 量化程序。例如，以下程序

func.func @add(%arg0: tensor<!quant.uniform<i8:f32,1.0:0>>, %arg1: tensor<!quant.uniform<i8:f32,2.0:1>>) ->  tensor<!quant.uniform<i8:f32,3.0:2>> {
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32,1.0:0>>, tensor<!quant.uniform<i8:f32,2.0:1>>) -> tensor<!quant.uniform<i8:f32,3.0:2>>
  func.return %0 : tensor<!quant.uniform<i8:f32,3.0:2>>
}

将变为：

func.func @add(%arg0: tensor<!quant.uniform<i8:f32, 1.000000e+00>>, %arg1: tensor<!quant.uniform<i8:f32, 2.000000e+00:1>>) -> tensor<!quant.uniform<i8:f32, 3.000000e+00:2>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<!quant.uniform<i8:f32, 1.000000e+00>>) -> tensor<f32>
  %1 = stablehlo.uniform_dequantize %arg1 : (tensor<!quant.uniform<i8:f32, 2.000000e+00:1>>) -> tensor<f32>
  %2 = stablehlo.add %0, %1 : tensor<f32>
  %3 = stablehlo.uniform_quantize %2 : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 3.000000e+00:2>>
  return %3 : tensor<!quant.uniform<i8:f32, 3.000000e+00:2>>
}

`-stablehlo-legalize-to-vhlo`

将 StableHLO 合法化为 VHLO。

将 StableHLO 合法化为 VHLO 中最新版本的操作。然后，可以使用 VhloToVersionPass 将这些操作降级为旧版 VHLO，以实现向前兼容性。

stablehlo.exponential %[[ARG0]] <{result_accuracy = DEFAULT}> : tensor<f32>
# ====>
"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.DEFAULT_v1}> : !vhlo.tensor_v1<!vhlo.f32_v1>

如需详细了解如何使用 VHLO 来保持向前和向后兼容性，请参阅 vhlo.md > VHLO 方言。

选项

-allow-other-dialects : Allow serialization to use other (potentially unstable) dialects, inserts unrealized casts between dialects.

`-stablehlo-refine-arguments`

优化了主函数的实参形状。

使用输入类型签名修改主函数的实参。将实参封装在 custom_call @stablehlo.shape_refinement_operand_wrapper 中，以在运行形状细化之前保持 IR 有效。

func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  ...
}

==>

func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {...}
    : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
}

refinedTypesOption 可用于指定一组精细类型。这可以在 MLIR 中使用 --types='tensor<...>,tensor<...>' 指定，也可以传递给 pass 创建方法。细化类型列表必须指定要细化的 main 方法的每个实参的类型。

选项

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

`-stablehlo-refine-shapes`

优化整个 StableHLO 程序的形状。

逐步介绍如何在操作中优化 StableHLO 程序的形状。

此 pass 的主要用途是将动态形状的程序专门化为静态形状。如果动态形状的 StableHLO 计划具有正确的结构，那么将其实参类型从动态形状更新为静态形状并运行此传递将会在整个计划中传播静态形状。

此传递通过直接将结果的使用替换为操作数来移除 custom_call @shape_refinement_operand_wrapper，并在整个程序中传播静态形状。

  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {...}
      : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  %1 = stablehlo.add %0, %0 : tensor<?xf32>

  ==>

  %1 = stablehlo.add %arg0, %arg0 : tensor<16xf32>

适用于形状细化的模块必须具有以下属性：

所有动态形状都仅依赖于输入形状（不依赖于输入数组内容）。我们将仅以传递方式依赖于输入形状（例如，由 stablehlo.get_dimension_size 给定）或全局常量（例如符号整数的已解析值，即张量：A = 5）的操作称为 dimension 操作。所有维度值都可以通过过程间常量折叠解析为常量。
中间函数可以在实参列表的开头采用多个令牌实参（类型为 !stablehlo.token），后面跟一些全局常量实参（即常量整数标量），例如符号整数的解析值（即张量：A = 5）。
某些中间函数可能会返回对全局常量的计算结果，即对 symint 值的 floordiv。这些函数仅在细化后返回常量值。这些函数是内联函数。
对单个函数的所有调用都会解析为相同的实参形状，并且不会进行任何递归 / 共同递归函数调用。

`-stablehlo-wrap-in-composite`

将非复合 StableHLO 操作封装在复合操作中。

将 StableHLO 操作封装在 stablehlo.composite 操作中。

例如，考虑一个简单的 StableHLO 程序：

func.func @main(%arg0 : tensor<2xf32>, %arg1 : tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
  return %0 : tensor<2xf32>
}

将此 pass 应用于封装 stablehlo.add 操作将生成以下程序：

func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.composite "stablehlo.add" %arg0, %arg1 {decomposition = @stablehlo.add.impl} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  return %0 : tensor<2xf32>
}
func.func private @stablehlo.add.impl(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
  return %0 : tensor<2xf32>
}

注意：

生成的 stablehlo.composite 操作的 name 属性将始终与被封装的原始操作的名称相同（例如，如果您封装 stablehlo.add 操作，复合操作也将命名为 "stablehlo.add"）。
封装原始操作（由 stablehlo.composite 操作的 decomposition 属性引用）的私有函数将使用 <op_name>.impl[.N] 模式命名，其中 <op_name> 是原始操作的名称，N 是生成的唯一整数标识符，用于防止模块内出现命名冲突。

此卡券可用于以下两种不同的方式：

模式 1：命令行用法

此模式适用于调试或测试，因为它可以对生成的 stablehlo.composite 操作的属性进行最少的控制。它会封装使用 op-names（以英文逗号分隔的操作名称列表）选项指定的所有操作实例。新创建的 stablehlo.composite 操作的属性将与原始操作的属性相同。

用法示例：

stablehlo-opt input.mlir --stablehlo-wrap-in-composite=op-names='stablehlo.add,stablehlo.mul' -o output.mlir

模式 2：通过自定义属性处理实现程序化模块级封装

此模式将程序化封装扩展到整个模块，可对封装的操作及其属性进行精细控制。这是通过使用 createStablehloWrapInCompositePass API 实现的，该 API 将 CompositeAttributeProviderMap 作为实参。

CompositeAttributeProviderMap 是一个映射，用于指定哪些操作应考虑进行封装以及应如何处理这些操作的属性。其语义如下：

键 (mlir::TypeID)：MLIR 操作的 TypeID。如果操作的 TypeID 与映射中的某个键匹配，则该操作会成为封装候选对象。
值（Lambda 函数）：类型为 std::function<std::optional<NamedAttrList>(Operation*)> 的 Lambda 函数。此函数会应用于每个候选操作。
- 输入：一个 mlir::Operation*，即与 TypeID 键对应的操作类型的实例。
- 返回值：一个 std::optional<NamedAttrList>。
  - 如果 lambda 返回 NamedAttrList（封装在 std::optional 中），则该操作封装在 stablehlo::composite 操作中，并且返回的属性用于设置复合的属性。
  - 如果 lambda 返回 std::nullopt，则操作不会被封装。这样就可以根据自定义条件进行选择性封装。

示例 (C++)：


stablehlo::CompositeAttributeProviderMap compositeAttributeProviderMap;

compositeAttributeProviderMap[mlir::TypeID::get<mlir::stablehlo::AddOp>()] =
  [](mlir::Operation* op) -> std::optional<mlir::NamedAttrList> {
  // Custom logic to determine if and how to wrap the operation.
  // Example: Only wrap if it's on a specific type.
  if (mlir::isa<mlir::Float32Type>(op->getOperand(0).getType())) {
    return mlir::NamedAttrList(op->getAttrs());
  }
  return std::nullopt; // Do not wrap.
};

pm.addPass(createStablehloWrapInCompositePass(compositeAttributeProviderMap, compositeVersion));
if (mlir::failed(pm.run(module))) {
  return;
}

选项

-op-names : The names of the ops to wrap.
-version  : The version number of the composite op.

`-vhlo-legalize-to-stablehlo`

将 VHLO 合法化为 StableHLO。

`-vhlo-to-version`

在不同版本的 VHLO 之间进行转换，以实现兼容性。

在 VHLO 版本之间进行转换，以实现 IR 升级和降级，从而保持向前和向后兼容性。

"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = DEFAULT}>
# ==( -target=1.0.0 )==>
"vhlo.exponential_v1"(%[[ARG0]])
# ==( -target=1.9.0 )==>
"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = DEFAULT}>

如需详细了解如何使用 VHLO 来保持向前和向后兼容性，请参阅 vhlo.md > VHLO 方言。

选项

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