StableHLO 规范

StableHLO 是机器学习 (ML) 模型中高级操作 (HLO) 的一组操作。StableHLO 可作为不同机器学习框架和机器学习编译器之间的可移植性层：生成 StableHLO 程序的机器学习框架与使用 StableHLO 程序的机器学习编译器兼容。

我们的目标是，通过在各种机器学习框架（例如 TensorFlow、JAX 和 PyTorch）和机器学习编译器（例如 XLA 和 IREE）之间创建更多互操作性，来简化和加速机器学习开发。为此，本文档提供了 StableHLO 编程语言的规范。

此规范包含三个主要部分。首先，程序部分介绍了 StableHLO 程序的结构，该程序由 StableHLO 函数组成，而 StableHLO 函数本身又由 StableHLO 操作组成。在该结构中，Ops 部分指定了各个操作的语义。执行部分提供了在程序中一起执行的所有这些操作的语义。最后，“表示法”部分讨论了整个规范中使用的表示法。

如需查看之前版本的 StableHLO 的规范，请打开感兴趣的已标记的版本的代码库。例如，StableHLO v0.19.0 规范。如需查看 StableHLO 在每次次要版本升级时发生的更改，请参阅 VhloDialect.td 中的版本日志。

计划

Program ::= {Func}

StableHLO 程序由任意数量的 StableHLO 函数组成。以下是一个示例程序，其中包含一个函数 @main，该函数有 3 个输入（%image、%weights 和 %bias）和 1 个输出。函数的正文有 6 个操作。

func.func @main(
  %image: tensor<28x28xf32>,
  %weights: tensor<784x10xf32>,
  %bias: tensor<1x10xf32>
) -> tensor<1x10xf32> {
  %0 = "stablehlo.reshape"(%image) : (tensor<28x28xf32>) -> tensor<1x784xf32>
  %1 = "stablehlo.dot"(%0, %weights) : (tensor<1x784xf32>, tensor<784x10xf32>) -> tensor<1x10xf32>
  %2 = "stablehlo.add"(%1, %bias) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  %3 = "stablehlo.constant"() {value = dense<0.0> : tensor<1x10xf32>} : () -> tensor<1x10xf32>
  %4 = "stablehlo.maximum"(%2, %3) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  "func.return"(%4): (tensor<1x10xf32>) -> ()
}

函数

Func        ::= 'func' '.' 'func' FuncId FuncInputs FuncOutputs '{' FuncBody '}'
FuncInputs  ::= '(' [FuncInput {',' FuncInput}] `)`
FuncInput   ::= ValueId ':' ValueType
FuncOutputs ::= ['->' FuncOutput, {',' FuncOutput}]
FuncOutput  ::= ValueType
FuncBody    ::= {Op}

StableHLO 函数（也称为命名函数）具有标识符、输入/输出和正文。未来，我们计划为函数引入更多元数据，以实现与 HLO 的更好兼容性（#425、#626、#740、#744）。

标识符

FuncId  ::= '@' letter {letter | digit}
ValueId ::= '%' digit {digit}
          | '%' letter {letter | digit}
letter  ::= 'a' | ... | 'z' | 'A' | ... | 'Z' | '_'
digit   ::= '0' | ... | '9'

StableHLO 标识符与许多编程语言中的标识符类似，但有两个特殊之处：1) 所有标识符都有标记，用于区分不同类型的标识符；2) 值标识符可以是纯数字，以简化 StableHLO 程序的生成。

类型

Type         ::= ValueType | NonValueType
ValueType    ::= TensorType | QuantizedTensorType | TokenType | TupleType | BufferType
NonValueType ::= TensorElementType | QuantizedTensorElementType | FunctionType | StringType

StableHLO 类型分为值类型（也称为一等类型），表示 StableHLO 值；以及非值类型，用于描述其他程序元素。StableHLO 类型与许多编程语言中的类型类似，主要特点是 StableHLO 的领域特定性质，这会导致一些不寻常的结果（例如，标量类型不是值类型）。

TensorType ::= 'tensor' '<' Shape TensorElementType '>'
Shape ::= {DimensionSize 'x'}
DimensionSize ::= digit {digit} | '?'

张量类型表示张量，即多维数组。它们具有形状和元素类型，其中形状表示非负或未知的维度大小，按相应维度（也称为轴）的升序排列，编号从 0 到 R-1。维度数 R 称为秩。例如，tensor<2x3xf32> 是一个形状为 2x3、元素类型为 f32 的张量类型。它有两个维度（或者说两个轴）- 第 0 个维度和第 1 个维度 - 大小分别为 2 和 3。其等级为 2。

形状可以是部分未知或完全未知（动态），例如 tensor<?x2xf64> 是部分未知，而 tensor<?x?xf64> 是完全未知。动态维度大小使用 ? 表示。无法取消对形状的排名。

未来，我们计划探索将张量类型扩展到维度大小和元素类型之外，例如，纳入布局 (#629) 和稀疏性 (#1078)。

QuantizedTensorType ::= 'tensor' '<' Shape QuantizedTensorElementType '>'
QuantizedTensorElementType ::= '!quant.uniform' '<'
                  QuantizationStorageType
                  ['<' QuantizationStorageMin ':' QuantizationStorageMax '>']
                  ':' QuantizationExpressedType
                  [':' QuantizationDimension]
                  ',' QuantizationParameters '>'
QuantizationStorageType ::= IntegerType
QuantizationStorageMin ::= IntegerLiteral
QuantizationStorageMax ::= IntegerLiteral
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerLiteral
QuantizationParameters ::= QuantizationParameter
                         | '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale [':' QuantizationZeroPoint]
QuantizationScale ::= FloatLiteral
QuantizationZeroPoint ::= IntegerLiteral

量化元素类型表示存储类型的整数值，范围从 storage_min 到 storage_max（含），对应于表达类型的浮点值。对于给定的整数值 i，相应的浮点值 f 可计算为 f = (i - zero_point) * scale，其中 scale 和 zero_point 称为量化参数。在语法中，storage_min 和 storage_max 是可选的，但默认值分别为 min_value(storage_type) 和 max_value(storage_type)。量化元素类型具有以下限制：

• (C1) type(storage_min) = storage_type。
• (C2) type(storage_max) = storage_type。
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type)。
• (C4) type(scales...) = expressed_type。
• (C5) 0 < scales。
• (C6) is_finite(scales...)。
• (C7) storage_min <= zero_points <= storage_max。
• (C8) type(zero_points...) = storage_type。
• (C9) size(scales) = size(zero_points)。
• (C10) 如果 is_empty(quantization_dimension)，则 size(scales) = 1。
• (C11) 0 <= quantization_dimension。

目前，QuantizationScale 是一个浮点常量，但人们对基于整数的缩放（用乘数和偏移量表示）非常感兴趣。我们计划在不久的将来探索此功能 (#1404)。

目前，我们正在讨论 QuantizationZeroPoint 的语义，包括类型、值以及量化张量类型中是否只能有一个或可能存在多个零点。根据此次讨论的结果，未来可能会更改有关零分的相关规范 (#1405)。

另一项正在进行的讨论涉及 QuantizationStorageMin 和 QuantizationStorageMax 的语义，以确定是否应针对这些值和量化张量的值施加任何限制 (#1406)。

最后，我们计划探索如何表示未知的比例和零点，这与我们计划探索如何表示未知的维度大小 (#1407) 类似。

量化张量类型表示具有量化元素的张量。这些张量与常规张量完全相同，只是其元素具有量化元素类型，而不是常规元素类型。

在量化张量中，量化可以是按张量进行的，即整个张量只有一个 scale 和 zero_point；也可以是按轴进行的，即有多个 scales 和 zero_points，每个特定维度 quantization_dimension 的切片都有一对。更正式地说，在具有按轴量化的张量 t 中，quantization_dimension 有 dim(t, quantization_dimension) 个切片：t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] 等。第 i 个切片中的所有元素都使用 scales[i] 和 zero_points[i] 作为其量化参数。量化张量类型具有以下限制：

• 对于按张量进行量化：
  • 没有其他限制。
• 对于按轴量化：
  • (C12) quantization_dimension < rank(self)。
  • (C13) dim(self, quantization_dimension) = size(scales)。

TokenType ::= 'token'

令牌类型表示令牌，即由某些操作生成和使用的不透明值。令牌用于对操作强制执行顺序，如执行部分中所述。

TupleType ::= 'tuple' '<' TupleElementTypes '>'
TupleElementTypes ::= [ValueType {',' ValueType}]

缓冲区类型表示缓冲区。例如，在 XLA 中，缓冲区是具有一致存储的多维数组。与张量类型类似，缓冲区类型具有形状和元素类型，其中形状表示非负或未知的维度大小，按相应维度（也称为轴）的升序排列，编号从 0 到 R-1。维度数 R 称为秩。例如，memref<2x3xf32> 是一个形状为 2x3、元素类型为 f32 的缓冲区类型。它有两个维度（或者说两个轴）- 第 0 个维度和第 1 个维度 - 大小分别为 2 和 3。其等级为 2。

可以使用 custom_call 到 CreateBuffer 或 Pin 分配缓冲区，并通过 custom_call 到 Unpin 取消分配缓冲区。只有 custom_call 操作可以读取和写入缓冲区内的内容。如需了解详情，请参阅 custom_call。

元组类型表示元组，即异构列表。元组是一项旧版功能，仅用于与 HLO 兼容。在 HLO 中，元组用于表示可变输入和输出。在 StableHLO 中，可变数量的输入和输出受到原生支持，并且 StableHLO 中元组的唯一用途是全面表示 HLO ABI，其中例如 T、tuple<T> 和 tuple<tuple<T>> 可能因特定实现而有很大差异。未来，我们计划对 HLO ABI 进行更改，这可能使我们能够从 StableHLO 中移除元组类型 (#598)。

TensorElementType ::= BooleanType | IntegerType | FloatType | ComplexType
BooleanType ::= 'i1'
IntegerType ::= SignedIntegerType | UnsignedIntegerType
SignedIntegerType ::= 'si2' | 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui2' | 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f4E2M1FN' | 'f6E2M3FN' | 'f6E3M2FN' | 'f8E3M4' | 'f8E4M3'
            | 'f8E4M3FN' | 'f8E4M3FNUZ' | 'f8E4M3B11FNUZ' | 'f8E5M2'
            | 'f8E5M2FNUZ' | 'f8E8M0FNU' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

元素类型表示张量类型的元素。与许多编程语言不同，这些类型在 StableHLO 中不是一等类型。这意味着 StableHLO 程序无法直接表示这些类型的值（因此，惯例是使用类型为 tensor<T> 的 0 维张量值来表示类型为 T 的标量值）。

• 布尔值类型表示布尔值 true 和 false。
• 整数类型可以是有符号 (si) 或无符号 (ui)，并且具有支持的位宽（2、4、8、16、32 或 64）。有符号 siN 类型表示从 -2^(N-1) 到 2^(N-1)-1（含）的整数值，无符号 uiN 类型表示从 0 到 2^N-1（含）的整数值。
• 浮点类型可以是以下类型之一：
  • f8E3M4、f8E4M3 和 f8E5M2 遵循 IEEE-754 惯例的 8 位浮点数。
  • f8E4M3FN 和 f8E5M2 类型分别对应于 FP8 Formats for Deep Learning 中描述的 FP8 格式的 E4M3 和 E5M2 编码。
  • f8E4M3FNUZ 和 f8E5M2FNUZ 类型分别对应于用于深度神经网络的 8 位数值格式中所述 FP8 格式的 E4M3 和 E5M2 编码。
  • f8E4M3B11FNUZ 类型，对应于用于深度神经网络的混合 8 位浮点 (HFP8) 训练和推理中所述的 FP8 格式的 E4M3 编码。
  • 与 BFloat16：Cloud TPU 上实现高性能的秘诀中所述的 bfloat16 格式对应的 bf16 类型。
  • f16、f32 和 f64 类型分别对应于 IEEE 754 标准中描述的 binary16（“半精度”）、binary32（“单精度”）和 binary64（“双精度”）格式。
  • tf32 类型对应于 TensorFloat32 格式，在 StableHLO 中仅受到有限的支持。
  • OCP 微缩格式规范中描述的 f4E2M1FN、f6E2M3FN、f6E3M2FN 和 f8E8M0FNU MX（微缩）类型。
• 复数类型表示具有相同元素类型的实部和虚部的复数值。支持的复杂类型包括 complex<f32>（两个部分均为 f32 类型）和 complex<f64>（两个部分均为 f64 类型）。

FunctionType ::= '(' InputTypes ')' '->' '(' OutputTypes ')'
InputTypes ::= [ValueType {',' ValueType}]
OutputTypes ::= [ValueType {',' ValueType}]

函数类型表示命名函数和匿名函数。它们具有输入类型（-> 左侧的类型列表）和输出类型（-> 右侧的类型列表）。在许多编程语言中，函数类型是一等类型，但在 StableHLO 中不是。

StringType ::= 'string'

字符串类型表示字节序列。与许多编程语言不同，字符串类型在 StableHLO 中不是一等公民，仅用于为程序元素指定静态元数据。

运维

StableHLO 操作（也称为 op）表示机器学习模型中的一组封闭的高级操作。如上所述，StableHLO 语法在很大程度上受到 MLIR 的启发，虽然 MLIR 不一定是最符合人体工程学的替代方案，但可以说是最适合 StableHLO 的，因为 StableHLO 的目标是在机器学习框架和机器学习编译器之间创建更好的互操作性。

Op            ::= [OpOutputs] OpName OpInputs ':' OpSignature
OpName        ::= '"' 'stablehlo' '.' OpMnemonic '"'
OpMnemonic    ::= 'abs' | 'add' | ...

StableHLO 操作（也称为 op）具有名称、输入/输出和签名。该名称由 stablehlo. 前缀和可唯一标识某个受支持操作的助记符 组成。如需查看所有受支持的运算的完整列表，请参阅下文。

OpInputs        ::= OpInputValues OpInputFuncs OpInputAttrs
OpInputValues   ::= '(' [OpInputValue {',' OpInputValue}] ')'
OpInputValue    ::= ValueId
OpInputFuncs    ::= ['(' OpInputFunc {',' OpInputFunc} ')']
OpInputAttrs    ::= ['{' OpInputAttr {',' OpInputAttr} '}']
OpOutputs       ::= [OpOutput {',' OpOutput} '=']
OpOutput        ::= ValueId

操作会使用输入并生成输出。输入分为输入值（在执行期间计算）、输入函数（静态提供，因为在 StableHLO 中，函数不是一级值）和输入属性（也是静态提供）。操作所消耗和产生的输入和输出类型取决于其助记符。例如，add op 会消耗 2 个输入值并生成 1 个输出值。相比之下，select_and_scatter 操作会消耗 3 个输入值、2 个输入函数和 3 个输入属性。

OpInputFunc ::= '{' Unused FuncInputs ':' FuncBody '}'
Unused      ::= '^' digit {digit}
              | '^' letter {letter | digit}

输入函数（也称为匿名函数）与命名函数非常相似，但有以下区别：1) 它们没有标识符（因此称为“匿名”），2) 它们不声明输出类型（输出类型是从函数中的 return 操作推断出来的）。

输入函数的语法包含一个目前未使用的部分（请参阅上面的 Unused 产品），该部分用于与 MLIR 兼容。在 MLIR 中，有一个更通用的“区域”概念，该区域可以包含多个通过跳转操作连接在一起的操作“块”。这些块具有与 Unused 产品对应的 ID，以便彼此区分。 StableHLO 没有跳转操作，因此 MLIR 语法的相应部分未使用（但仍然存在）。

OpInputAttr      ::= OpInputAttrName '=' OpInputAttrValue
OpInputAttrName  ::= letter {letter | digit}
OpInputAttrValue ::= Constant

输入属性具有名称和值，该值是受支持的常量之一。它们是为程序元素指定静态元数据的主要方式。例如，concatenate 操作使用属性 dimension 来指定串联输入值的维度。同样，slice 操作使用多个属性（例如 start_indices 和 limit_indices）来指定用于对输入值进行切片的边界。

目前，实际应用中的 StableHLO 程序有时包含本文档中未描述的属性。未来，我们计划将这些属性纳入 StableHLO opset，或者禁止它们出现在 StableHLO 程序中。在此期间，您可以查看以下属性列表：

• layout (#629)。
• mhlo.frontend_attributes (#628)。
• mhlo.sharding (#619)。
• output_operand_aliases (#740)。
• 位置元数据 (#594)。

OpSignature ::= '(' [ValueType {',' ValueType}] ')' '->' '(' [ValueType {',' ValueType}] ')'

操作签名包含所有输入值的类型（-> 左侧的类型列表）和所有输出值的类型（-> 右侧的类型列表）。严格来说，输入类型是冗余的，输出类型也几乎总是冗余的（因为对于大多数 StableHLO 操作，输出类型可以从输入中推断出来）。不过，为了与 MLIR 兼容，op 签名特意纳入了 StableHLO 语法。

以下是一个助记符为 select_and_scatter 的操作示例。它使用 3 个输入值（%operand、%source 和 %init_value）、2 个输入函数和 3 个输入属性（window_dimensions、window_strides 和 padding）。请注意，该操作的签名仅包含其输入值的类型（但不包含以内联方式提供的输入函数和属性的类型）。

%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "stablehlo.return"(%0) : (tensor<i32>) -> ()
}) {
  window_dimensions = dense<[3, 1]> : tensor<2xi64>,
  window_strides = dense<[2, 1]> : tensor<2xi64>,
  padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi32>, tensor<2x2xi32>, tensor<i32>) -> tensor<4x2xi32>

常量

Constant ::= BooleanConstant
           | IntegerConstant
           | FloatConstant
           | ComplexConstant
           | TensorConstant
           | QuantizedTensorConstant
           | StringConstant
           | EnumConstant

StableHLO 常量具有字面量和类型，两者共同表示一个 StableHLO 值。一般来说，类型是常量语法的一部分，除非类型明确无误（例如，布尔常量明确具有 i1 类型，而整数常量可能具有多种类型）。

BooleanConstant ::= BooleanLiteral
BooleanLiteral  ::= 'true' | 'false'

布尔常量表示布尔值 true 和 false。布尔值常量的类型为 i1。

IntegerConstant   ::= IntegerLiteral ':' IntegerType
IntegerLiteral    ::= ['-' | '+'] DecimalDigits
                    | ['-' | '+'] '0x' HexadecimalDigits
DecimalDigits     ::= decimalDigit {decimalDigit}
HexadecimalDigits ::= hexadecimalDigit {hexadecimalDigit}
decimalDigit      ::= '0' | ... | '9'
hexadecimalDigit  ::= decimalDigit | 'a' | ... | 'f' | 'A' | ... | 'F'

整数常量通过使用十进制或十六进制表示法的字符串来表示整数值。不支持其他进制，例如二进制或八进制。 整数常量具有以下限制：

• (C1) is_wellformed(integer_literal, integer_type)。

FloatConstant  ::= FloatLiteral ':' FloatType
FloatLiteral   ::= SignPart IntegerPart FractionalPart ScientificPart
                 | '0x' [HexadecimalDigits]
SignPart       ::= ['-' | '+']
IntegerPart    ::= DecimalDigits
FractionalPart ::= ['.' [DecimalDigits]]
ScientificPart ::= [('e' | 'E') ['-' | '+'] DecimalDigits]

浮点常量通过使用十进制或科学记数法的字符串表示浮点值。此外，十六进制表示法可用于直接指定相应类型的浮点格式中的底层位。浮点常量具有以下限制：

• (C1) 如果使用非十六进制表示法，则为 is_wellformed(float_literal, float_type)。
• (C2) 如果使用十六进制表示法，则为 size(hexadecimal_digits) = num_bits(float_type) / 4。

ComplexConstant ::= ComplexLiteral ':' ComplexType
ComplexLiteral  ::= '(' RealPart ',' ImaginaryPart ')'
RealPart        ::= FloatLiteral
ImaginaryPart   ::= FloatLiteral

复数常量使用实部（在前）和虚部（在后）的列表来表示复数值。例如，(1.0, 0.0) : complex<f32> 表示 1.0 + 0.0i，(0.0, 1.0) : complex<f32> 表示 0.0 + 1.0i。然后，这些部分在内存中的存储顺序由实现定义。复杂常数具有以下限制：

• (C1) is_wellformed(real_part, complex_element_type(complex_type))。
• (C2) is_wellformed(imaginary_part, complex_element_type(complex_type))。

TensorConstant ::= TensorLiteral ':' TensorType
TensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'
DenseLiteral   ::= DenseDimension | DenseElements
DenseDimension ::= '[' [DenseLiteral {',' DenseLiteral}] ']'
DenseElements  ::= [ElementLiteral {',' ElementLiteral}]
ElementLiteral ::= BooleanLiteral | IntegerLiteral | FloatLiteral | ComplexLiteral

张量常量使用通过 NumPy 表示法指定的嵌套列表来表示张量值。例如，dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> 表示一个张量值，其从索引到元素的映射关系如下：{0, 0} => 1、{0, 1} => 2、{0, 2} => 3、{1, 0} => 4、{1, 1} => 5、{1, 2} => 6。然后，这些元素在内存中的存储顺序由实现定义。张量常量具有以下限制：

• (C1) has_syntax(tensor_literal, element_type(tensor_type))，其中：
  • has_syntax(element_literal: Syntax, element_type: Type) = is_wellformed(element_literal, type)。
  • has_syntax(tensor_literal: List, element_type: Type) = has_syntax(tensor_literal..., element_type)。
• (C2) has_shape(tensor_literal, shape(tensor_type))，其中：
  • has_shape(element_literal: Syntax, []) = true。
  • has_shape(tensor_literal: List, shape: List) = size(tensor_literal) = shape[0] and has_shape(tensor_literal..., shape[1:])。
  • 否则，false。

QuantizedTensorConstant ::= QuantizedTensorLiteral ':' QuantizedTensorType
QuantizedTensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'

量化张量常量使用与张量常量相同的表示法来表示量化张量值，元素指定为其存储类型的常量。量化张量常量具有以下限制：

• (C1) has_syntax(quantized_tensor_literal, storage_type(quantized_tensor_type))。
• (C2) has_shape(quantized_tensor_literal, shape(quantized_tensor_type))。

StringConstant  ::= StringLiteral
StringLiteral   ::= '"' {stringCharacter | escapeSequence} '"'
stringCharacter ::= all ASCII characters except '\00', '\01', ... '\1f' and '"'
escapeSequence  ::= '\' ('"' | '\' | 'n' | 't' | (hexadecimalDigit hexadecimalDigit))

字符串字面值由使用 ASCII 字符和转义序列指定的字节组成。它们与编码无关，因此这些字节的解释由实现定义。字符串字面量的类型为 string。

操作/运算

abs

语义

对 operand 张量执行元素级绝对值运算，并生成 result 张量。根据元素类型，执行以下操作：

• 对于有符号整数：整数模数。
• 对于浮点数：来自 IEEE-754 的 abs。
• 对于复数：复数模。
• 对于量化类型：dequantize_op_quantize(abs, operand, type(result))。

输入

输出

限制条件

• (C1) shape(result) = shape(operand)。
• (C2) baseline_element_type(result) 的定义如下：
  • 如果 is_complex(operand)，则为 complex_element_type(element_type(operand))。
  • 否则为 baseline_element_type(operand)。

示例

// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand) : (tensor<3xi32>) -> tensor<3xi32>
// %result: [2, 0, 2]

更多示例

添加

语义

对两个张量 lhs 和 rhs 执行逐元素加法，并生成一个 result 张量。根据元素类型，执行以下操作：

• 对于布尔值：逻辑或。
• 对于整数：整数加法。
• 对于浮点数：来自 IEEE-754 的 addition。
• 对于复数：复数加法。
• 对于量化类型：dequantize_op_quantize(add, lhs, rhs, type(result))。

输入

输出

限制条件

• 如果运算使用非量化张量：
  • (C1) type(lhs) = type(rhs) = type(result)。
• 如果操作使用量化张量：
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result)。
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result)。
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)。
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result)。
  • (C6) 如果 is_per_axis_quantized(lhs)，则 quantization_dimension(lhs) = quantization_dimension(result)。
  • (C7) 如果 is_per_axis_quantized(rhs)，则 quantization_dimension(rhs) = quantization_dimension(result)。

示例

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.add"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[6, 8], [10, 12]]

更多示例

after_all

语义

确保生成 inputs 的操作在任何依赖于 result 的操作之前执行。执行此操作不会执行任何操作，它仅用于建立从 result 到 inputs 的数据依赖关系。

输入

输出

示例

// %input0: !stablehlo.token
// %input1: !stablehlo.token
%result = "stablehlo.after_all"(%input0, %input1) : (!stablehlo.token, !stablehlo.token) -> !stablehlo.token

更多示例

all_gather

语义

在 StableHLO 进程网格中的每个进程组内，沿 all_gather_dim 连接来自每个进程的 operands 张量的值，并生成 results 张量。

该操作将 StableHLO 进程网格拆分为 process_groups，定义如下：

• 如果 channel_id <= 0 and use_global_device_ids = false，则为 cross_replica(replica_groups)。
• 如果 channel_id > 0 and use_global_device_ids = false，则为 cross_replica_and_partition(replica_groups)。
• 如果 channel_id > 0 and use_global_device_ids = true，则为 flattened_ids(replica_groups)。

之后，在每个 process_group 中：

• operands...@receiver = [operand@sender for sender in process_group]（适用于process_group中的所有receiver）。
• results...@process = concatenate(operands...@process, all_gather_dim)（适用于process_group中的所有process）。

输入

输出

限制条件

• (C1) 0 <= all_gather_dim < rank(operands...)。
• (C2) is_unique(replica_groups)。
• (C3) size(replica_groups) 的定义如下：
  • 如果使用 cross_replica，则为 num_replicas。
  • 如果使用 cross_replica_and_partition，则为 num_replicas。
  • 如果使用 flattened_ids，则为 num_processes。
• (C4) 0 <= replica_groups < size(replica_groups)。
• (C5) 如果 use_global_device_ids = true，则 channel_id > 0。
• (C6) type(results...) = type(operands...) 除外：
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1)。

示例

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [[1, 2], [3, 4]]
// %operand0@(1, 0): [[5, 6], [7, 8]]
// %operand1@(0, 0): [[11, 12], [13, 14]]
// %operand1@(1, 0): [[15, 16], [17, 18]]
%result:2 = "stablehlo.all_gather"(%operand0, %operand1) {
  all_gather_dim = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

更多示例

all_reduce

语义

在 StableHLO 进程网格中的每个进程组内，将缩减函数 computation 应用于每个进程的 operands 张量的值，并生成 results 张量。

该操作将 StableHLO 进程网格拆分为 process_groups，定义如下：

• 如果 channel_id <= 0 and use_global_device_ids = false，则为 cross_replica(replica_groups)。
• 如果 channel_id > 0 and use_global_device_ids = false，则为 cross_replica_and_partition(replica_groups)。
• 如果 channel_id > 0 and use_global_device_ids = true，则为 flattened_ids(replica_groups)。

之后，在每个 process_group 中：

• 对于某个二叉树 results...@process[result_index] = exec(schedule)，schedule 其中：
  • exec(node) = computation(exec(node.left), exec(node.right))。
  • exec(leaf) = leaf.value。
• schedule 是一个由实现定义的二叉树，其按顺序遍历为 to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0]))。

输入

输出

限制条件

• (C1) is_unique(replica_groups)。
• (C2) size(replica_groups) 的定义如下：
  • 如果使用 cross_replica，则为 num_replicas。
  • 如果使用 cross_replica_and_partition，则为 num_replicas。
  • 如果使用 flattened_ids，则为 num_processes。
• (C3) 0 <= replica_groups < size(replica_groups)。
• (C4) 如果 use_global_device_ids = true，则 channel_id > 0。
• (C5) computation 的类型为 (tensor<E>, tensor<E>) -> (tensor<E>)，其中 is_promotable(element_type(operand), E)。
• (C6) shape(results...) = shape(operands...)。
• (C7) element_type(results...) = E。

示例

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [1, 2, 3, 4]
// %operand0@(1, 0): [5, 6, 7, 8]
// %operand1@(0, 0): [9, 10, 11, 12]
// %operand1@(1, 0): [13, 14, 15, 16]
%result:2 = "stablehlo.all_reduce"(%operand0, %operand0) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

更多示例

all_to_all

语义

在 StableHLO 进程网格中的每个进程组内，沿 split_dimension 将 operands 张量的值拆分为若干部分，在进程之间分散拆分后的部分，沿 concat_dimension 连接分散的部分，并生成 results 张量。 该操作将 StableHLO 进程网格拆分为 process_groups，定义如下：

• 如果 channel_id <= 0，则为 cross_replica(replica_groups)。
• 如果 channel_id > 0，则为 cross_partition(replica_groups)。

之后，在每个 process_group 中：

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) 适用于 process_group 中的所有 sender。
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group]，其中receiver_index = process_group.index(receiver)。
• results...@process = concatenate(scattered_parts...@process, concat_dimension)。

输入

输出

限制条件

• (C1) 0 <= split_dimension < rank(operands...)。
• (C2) dim(operands..., split_dimension) % split_count = 0。
• (C3) 0 <= concat_dimension < rank(operands...)。
• (C4) 0 < split_count。
• (C5) is_unique(replica_groups)。
• (C6) size(replica_groups) 的定义如下：
  • 如果使用 cross_replica，则为 num_replicas。
  • 如果使用 cross_partition，则为 num_partitions。
• (C7) 0 <= replica_groups < size(replica_groups)。
• (C8) dim(replica_groups, 1) = split_count。
• (C9) type(results...) = type(operands...) 但如果 split_dimension != concat_dimension，则除外：
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count。
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count。

示例

// num_replicas: 2
// num_partitions: 1
// %operand1@(0, 0): [[1, 2, 3, 4],
//                    [5, 6, 7, 8]]
// %operand1@(1, 0): [[9, 10, 11, 12],
//                    [13, 14, 15, 16]]
// %operand2@(0, 0): [[17, 18, 19, 20],
//                    [21, 22, 23, 24]]
// %operand2@(1, 0): [[25, 26, 27, 28],
//                    [29, 30, 31, 32]]
%result:2 = "stablehlo.all_to_all"(%operand1, %operand2) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
} : (tensor<2x4xi64>, tensor<2x4xi64>) -> (tensor<4x2xi64>, tensor<4x2xi64>)
// %result#0@(0, 0): [[1, 2], [5, 6], [9, 10], [13, 14]]
// %result#0@(1, 0): [[3, 4], [7, 8], [11, 12], [15, 16]]
// %result#1@(0, 0): [[17, 18], [21, 22], [25, 26], [29, 30]]
// %result#1@(1, 0): [[19, 20], [23, 24], [27, 28], [31, 32]]

更多示例

和

语义

对两个张量 lhs 和 rhs 执行按元素与运算，并生成一个 result 张量。根据元素类型，执行以下操作：

• 对于布尔值：逻辑 AND。
• 对于整数：按位 AND。

输入

输出

限制条件

• (C1) type(lhs) = type(rhs) = type(result)。

示例

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.and"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[1, 2], [3, 0]]

更多示例

atan2

语义

对 lhs 和 rhs 张量执行逐元素 atan2 运算，并生成 result 张量。根据元素类型，执行以下操作：

• 对于浮点数：来自 IEEE-754 的 atan2。
• 对于复数：复数 atan2。
• 对于量化类型：dequantize_op_quantize(atan2, lhs, rhs, type(result))。

输入

输出

限制条件

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)。

示例

// %lhs: [0.0, 1.0, -1.0]
// %rhs: [0.0, 0.0, 0.0]
%result = "stablehlo.atan2"(%lhs, %rhs) : (tensor<3xf64>, tensor<3xf64>) -> tensor<3xf64>
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

更多示例

batch_norm_grad

语义

计算 batch_norm_training 的多个输入的梯度（从 grad_output 反向传播），并生成 grad_operand、grad_scale 和 grad_offset 张量。更正式地说，此操作可以使用 Python 语法表示为对现有 StableHLO 操作的分解，如下所示：

def compute_sum(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  return sum

def compute_mean(operand, feature_index):
  sum = compute_sum(operand, feature_index)
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index):
  # Broadcast inputs to type(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance`
  # Intermediate values will be useful for computing gradients
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)

  # Use the implementation from batchnorm_expander.cc in XLA
  # Temporary variables have exactly the same names as in the C++ code
  elements_per_feature = broadcast_in_dim(
      constant(divide(size(operand), dim(operand, feature_index)),
               element_type(grad_output)),
      [], type(operand))
  i1 = multiply(grad_output, elements_per_feature)
  i2 = broadcast_in_dim(
      compute_sum(grad_output, feature_index), [feature_index], type(operand))
  i3 = broadcast_in_dim(
      compute_sum(multiply(grad_output, centered_operand), feature_index),
      [feature_index], type(operand))
  i4 = multiply(i3, centered_operand)
  i5 = divide(i4, add(variance_bcast, epsilon_bcast))
  i6 = subtract(subtract(i1, i2), i5)

  grad_operand =
      multiply(divide(divide(scale_bcast, stddev), elements_per_feature), i6)
  grad_scale =
      compute_sum(multiply(grad_output, normalized_operand), feature_index)
  grad_offset = compute_sum(grad_output, feature_index)

  return grad_operand, grad_scale, grad_offset

对于量化类型，执行 dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index))。

输入

输出

限制条件

• (C1) 0 <= feature_index < rank(operand)。
• (C2) operand、scale、mean、variance、grad_output、grad_operand、grad_scale 和 grad_offset 具有相同的 baseline_element_type。
• (C3) operand、grad_output 和 grad_operand 具有相同的形状。
• (C4) scale、mean、variance、grad_scale 和 grad_offset 具有相同的形状。
• (C5) size(scale) = dim(operand, feature_index)。

示例

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
// %grad_output: [
//                [[0.1, 0.1], [0.1, 0.1]],
//                [[0.1, 0.1], [0.1, 0.1]]
//               ]
%grad_operand, %grad_scale, %grad_offset =
"stablehlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>,
     tensor<2x2x2xf64>) -> (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %grad_operand: [
//                 [[0.0, 0.0], [0.0, 0.0]],
//                 [[0.0, 0.0], [0.0, 0.0]]
//                ]
// %grad_scale:  [0.0, 0.0]
// %grad_offset: [0.4, 0.4]

batch_norm_inference

语义

对 operand 张量进行归一化处理，但 feature_index 维度除外，并生成 result 张量。更正式地说，此操作可以使用 Python 语法表示为对现有 StableHLO 操作的分解，如下所示：

def batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index):
  # Broadcast inputs to shape(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  offset_bcast = broadcast_in_dim(offset, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance` instead of
  # computing them like `batch_norm_training` does.
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)
  return add(multiply(scale_bcast, normalized_operand), offset_bcast)

对于量化类型，执行 dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result))。

输入

输出

限制条件

• (C1) 0 <= feature_index < rank(operand)。
• (C2) operand、scale、offset、mean、variance 和 result 具有相同的 baseline_element_type。
• (C3) size(scale) = dim(operand, feature_index)。
• (C4) size(offset) = dim(operand, feature_index)。
• (C5) size(mean) = dim(operand, feature_index)。
• (C6) size(variance) = dim(operand, feature_index)。
• (C7) baseline_type(operand) = baseline_type(result)。

示例

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
%result = "stablehlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>) -> tensor<2x2x2xf64>
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

batch_norm_training

语义

计算除 feature_index 维度之外的所有维度上的平均值和方差，并对 operand 张量进行归一化处理，生成 output、batch_mean 和 batch_var 张量。更正式地说，此操作可以使用 Python 语法分解为现有的 StableHLO 操作，如下所示：

def compute_mean(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def compute_variance(operand, feature_index):
  mean = compute_mean(operand, feature_index)
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  centered_operand = subtract(operand, mean_bcast)
  return compute_mean(mul(centered_operand, centered_operand), feature_index)

def batch_norm_training(operand, scale, offset, epsilon, feature_index):
  mean = compute_mean(operand, feature_index)
  variance = compute_variance(operand, feature_index)
  return batch_norm_inference(operand, scale, offset, mean, variance, epsilon,
                              feature_index),
         mean, variance

对于量化类型，执行 dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var))。

输入

输出

限制条件

• (C1) 0 <= feature_index < rank(operand)。
• (C2) operand、scale、offset、batch_mean、batch_var 和 output 具有相同的 baseline_element_type。
• (C3) size(scale) = dim(operand, feature_index)。
• (C4) size(offset) = dim(operand, feature_index)。
• (C5) size(batch_mean) = dim(operand, feature_index)。
• (C6) size(batch_var) = dim(operand, feature_index)。
• (C7) baseline_type(output) = baseline_type(operand)。

示例

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
%output, %batch_mean, %batch_var = "stablehlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>) ->
    (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %output: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]
// %batch_mean: [2.0, 3.0]
// %batch_var: [1.0, 1.0]

bitcast_convert

语义

对 operand 张量执行位转换操作，并生成一个 result 张量，其中整个 operand 张量的位使用 result 张量的类型重新解释。

更正式地说，给定 E = element_type(operand)、E' = element_type(result) 和 R = rank(operand)：

• 如果 num_bits(E') < num_bits(E)，则为 bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1])。
• 如果 num_bits(E') > num_bits(E)，则为 bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :])。
• 如果 num_bits(E') = num_bits(E)，则为 bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1])。

bits 会返回给定值的内存中表示形式，并且其行为由实现定义，因为张量的确切表示形式由实现定义，并且元素类型的确切表示形式也由实现定义。

输入

输出

限制条件

• (C1) 假定有 E = is_quantized(operand) ? storage_type(operand) : element_type(operand)、E' = is_quantized(result) ? storage_type(result) : element_type(result) 和 R = rank(operand)：
  • 如果为 num_bits(E') = num_bits(E)，则为 shape(result) = shape(operand)。
  • 如果 num_bits(E') < num_bits(E)：
  • rank(result) = R + 1。
  • dim(result, i) = dim(operand, i) 适用于所有 0 <= i < R。
  • dim(result, R) * num_bits(E') = num_bits(E)。
  • 如果 num_bits(E') > num_bits(E)：
  • rank(result) = R - 1。
  • dim(result, i) = dim(operand, i) 适用于所有 0 <= i < R。
  • dim(operand, R - 1) * num_bits(E) = num_bits(E')。
• (C2) 如果 is_complex(operand) or is_complex(result)，则 is_complex(operand) and is_complex(result)。

示例

// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand) : (tensor<f64>) -> tensor<4xf16>
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

更多示例

broadcast_in_dim

语义

通过复制 operand 张量中的数据来扩展输入张量的维度和/或秩，并生成 result 张量。更正式地说，result[result_index] = operand[operand_index]，其中对于 axes(operand) 中的所有 d：

• 如果 dim(operand, d) = 1，则为 operand_index[d] = 0。
• 否则为 operand_index[d] = result_index[broadcast_dimensions[d]]。

输入

输出

限制条件

• (C1) element_type(result) 由以下人员提供：
  • element_type(operand)（如果 !is_per_axis_quantized(operand)）。
  • element_type(operand)，但 quantization_dimension(operand)、scales(operand) 和 zero_points(operand) 可能分别与 quantization_dimension(result)、scales(result) 和 zero_points(result) 不同，否则。
• (C2) size(broadcast_dimensions) = rank(operand)。
• (C3) 0 <= broadcast_dimensions < rank(result)。
• (C4) is_unique(broadcast_dimensions)。
• (C5) 对于 axes(operand) 中的所有 d：
  • dim(operand, d) = 1 或
  • dim(operand, d) = dim(result, broadcast_dimensions[d])。
• (C6) 如果 is_per_axis_quantized(result)：
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)]。
  • 如果值为 dim(operand, quantization_dimension(operand)) = 1，则值为 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))。

示例

// %operand: [
//            [1, 2, 3]
//           ]
%result = "stablehlo.broadcast_in_dim"(%operand) {
  broadcast_dimensions = array<i64: 2, 1>
} : (tensor<1x3xi32>) -> tensor<2x3x2xi32>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

更多示例

场景

语义

根据 index 的值，通过执行 branches 中的一个函数来生成输出。更正式地说，result = selected_branch()，其中：

• 如果 0 <= index < size(branches)，则为 selected_branch = branches[index]。
• 否则为 selected_branch = branches[-1]。

输入

输出

限制条件

• (C1) 0 < size(branches)。
• (C2) input_types(branches...) = []。
• (C3) same(output_types(branches...))。
• (C4) type(results...) = output_types(branches[0])。

示例

// %index: -1
// %result_branch0: [0, 0]
// %result_branch1: [1, 1]
%result0, %result1 = "stablehlo.case"(%index) ({
  "stablehlo.return"(%result_branch0, %result_branch0) : (tensor<2xi64>, tensor<2xi64>) -> ()
}, {
  "stablehlo.return"(%result_branch1, %result_branch1) : (tensor<2xi64>, tensor<2xi64>) -> ()
}) : (tensor<i32>) -> (tensor<2xi64>, tensor<2xi64>)
// %result0: [1, 1]
// %result1: [1, 1]

更多示例

cbrt

语义

对 operand 张量执行逐元素立方根运算，并生成 result 张量。根据元素类型，执行以下操作：

• 对于浮点数：来自 IEEE-754 的 rootn(x, 3)。
• 对于复数：复数立方根。
• 对于量化类型：dequantize_op_quantize(cbrt, operand, type(result))

输入

输出

限制条件

• (C1) baseline_type(operand) = baseline_type(result)。

示例

// %operand: [0.0, 1.0, 8.0, 27.0]
%result = "stablehlo.cbrt"(%operand) : (tensor<4xf64>) -> tensor<4xf64>
// %result: [0.0, 1.0, 2.0, 3.0]

更多示例

ceil

语义

对 operand 张量执行逐元素向上取整运算，并生成 result 张量。 实现 IEEE-754 规范中的 roundToIntegralTowardPositive 操作。对于量化类型，执行 dequantize_op_quantize(ceil, operand, type(result))。

输入

输出

限制条件

• (C1) baseline_type(operand) = baseline_type(result)。

示例

// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.ceil"(%operand) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-0.0, -0.0, 1.0, 1.0, 2.0]

更多示例

cholesky

语义

计算一批矩阵的 Cholesky 分解。

更正式地说，对于 index_space(result) 中的所有 i，result[i0, ..., iR-3, :, :] 是 a[i0, ..., iR-3, :, :] 的 Cholesky 分解，形式为下三角矩阵（如果 lower 为 true）或上三角矩阵（如果 lower 为 false）。 相反三角形（即严格上三角或严格下三角）中的输出值由实现定义。

如果存在输入矩阵不是 Hermitian 正定矩阵的 i，则行为未定义。

对于量化类型，执行 dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result))。

输入

输出

限制条件

• (C1) baseline_type(a) = baseline_type(result)。
• (C2) 2 <= rank(a)。
• (C3) dim(a, -2) = dim(a, -1)。

示例

// %a: [
//      [1.0, 2.0, 3.0],
//      [2.0, 20.0, 26.0],
//      [3.0, 26.0, 70.0]
//     ]
%result = "stablehlo.cholesky"(%a) {
  lower = true
} : (tensor<3x3xf32>) -> tensor<3x3xf64>
// %result: [
//           [1.0, 0.0, 0.0],
//           [2.0, 4.0, 0.0],
//           [3.0, 5.0, 6.0]
//          ]

限制取值范围

语义

将 operand 张量的每个元素限制在最小值和最大值之间，并生成 result 张量。更正式地说，result[result_index] = minimum(maximum(operand[result_index], min_element), max_element)，其中 min_element = rank(min) = 0 ? min[] : min[result_index]，max_element = rank(max) = 0 ? max[] : max[result_index]。对于量化类型，执行 dequantize_op_quantize(clamp, min, operand, max, type(result))。

对复数施加排序会涉及令人惊讶的语义，因此我们计划在未来移除对该操作的复数支持 (#560)。

输入

输出

限制条件

• (C1) rank(min) = 0 or shape(min) = shape(operand)。
• (C2) rank(max) = 0 or shape(max) = shape(operand)。
• (C3) baseline_element_type(min) = baseline_element_type(operand) = baseline_element_type(max)。
• (C4) baseline_type(operand) = baseline_type(result)。

示例

// %min: [5, 10, 15]
// %operand: [3, 13, 23]
// %max: [10, 15, 20]
%result = "stablehlo.clamp"(%min, %operand, %max) : (tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32>
// %result: [5, 13, 20]

更多示例

collective_broadcast

语义

在 StableHLO 进程网格中的每个进程组内，将源进程中的 operand 张量的值发送到目标进程，并生成 result 张量。

该操作将 StableHLO 进程网格拆分为 process_groups，定义如下：

• 如果 channel_id <= 0，则为 cross_replica(replica_groups)。
• 如果 channel_id > 0，则为 cross_partition(replica_groups)。

之后，result@process 由以下公式给出：

• 如果存在 i 使得进程处于 process_groups[i] 中，则为 operand@process_groups[i, 0]。
• 否则为 broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result))。

输入

输出

限制条件

• (C1) is_unique(replica_groups)。
• (C2) 0 <= replica_groups < N，其中 N 定义如下：
  • 如果使用 cross_replica，则为 num_replicas。
  • 如果使用 cross_partition，则为 num_partitions。
• (C3) type(result) = type(operand)。

示例

// num_replicas: 4
// num_partitions: 1
// %operand@(0, 0): [[1, 2]]
// %operand@(1, 0): [[3, 4]]
// %operand@(2, 0): [[5, 6]]
// %operand@(3, 0): [[7, 8]]
%result = "stablehlo.collective_broadcast"(%operand) {
  replica_groups = dense<[[2, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor1x2xi64>) -> tensor<1x2xi64>
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[5, 6]]
// %result@(2, 0): [[5, 6]]
// %result@(3, 0): [[0, 0]]

collective_permute

语义

在 StableHLO 进程网格中的每个进程组内，将源进程中的 operand 张量的值发送到目标进程，并生成 result 张量。

该操作将 StableHLO 进程网格拆分为 process_groups，定义如下：

• 如果 channel_id <= 0，则为 cross_replica(source_target_pairs)。
• 如果 channel_id > 0，则为 cross_partition(source_target_pairs)。

之后，result@process 由以下公式给出：

• operand@process_groups[i, 0]，如果存在 i 使得 process_groups[i, 1] = process。
• 否则为 broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result))。

输入

输出

限制条件

• (C1) dim(source_target_pairs, 1) = 2。
• (C2) is_unique(source_target_pairs[:, 0])。
• (C3) is_unique(source_target_pairs[:, 1])。
• (C4) 0 <= source_target_pairs < N，其中 N 定义为：
  • 如果使用 cross_replica，则为 num_replicas。
  • 如果使用 cross_partition，则为 num_partitions。
• (C5) type(result) = type(operand)。

示例

// num_replicas: 3
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
// %operand@(2, 0): [[9, 10], [11, 12]]
%result = "stablehlo.collective_permute"(%operand) {
  source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<2x2xi64>) -> tensor<2x2xi64>
//
// %result@(0, 0): [[0, 0], [0, 0]]
// %result@(1, 0): [[1, 2], [3, 4]]
// %result@(2, 0): [[5, 6], [7, 8]]

更多示例

比较

语义

根据 comparison_direction 和 compare_type 对 lhs 和 rhs 张量执行逐元素比较，并生成 result 张量。

comparison_direction 和 compare_type 的值具有以下语义：

对于布尔值和整数元素类型：

• EQ：lhs = rhs。
• NE：lhs != rhs。
• GE：lhs >= rhs。
• GT：lhs > rhs。
• LE：lhs <= rhs。
• LT：lhs < rhs。

对于具有 compare_type = FLOAT 的浮点元素类型，相应操作会实现以下 IEEE-754 操作：

• EQ：compareQuietEqual。
• NE：compareQuietNotEqual。
• GE：compareQuietGreaterEqual。
• GT：compareQuietGreater。
• LE：compareQuietLessEqual。
• LT：compareQuietLess。

对于具有 compare_type = TOTALORDER 的浮点元素类型，相应操作使用 IEEE-754 中的 totalOrder 和 compareQuietEqual 操作的组合。

对于复杂元素类型，系统会使用提供的 comparison_direction 和 compare_type 对 (real, imag) 对进行字典序比较。对复数施加排序会涉及令人惊讶的语义，因此我们计划在未来移除当 comparison_direction 为 GE、GT、LE 或 LT 时的复数支持 (#560)。

对于量化类型，执行 dequantize_compare(lhs, rhs, comparison_direction)。

输入

输出

限制条件

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs)。
• (C2) shape(lhs) = shape(rhs) = shape(result)。
• (C3) compare_type 的定义如下：
  • 如果 is_signed_integer(element_type(lhs))，则为 SIGNED。
  • 如果 is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs))，则为 UNSIGNED。
  • 如果 is_float(element_type(lhs))，则为 FLOAT 或 TOTALORDER。
  • 如果 is_complex(element_type(lhs))，则为 FLOAT。

示例

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = #stablehlo<comparison_direction LT>,
  compare_type = #stablehlo<comparison_type FLOAT>
} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>
// %result: [true, false]

更多示例

复杂

语义

根据实数值和虚数值对 lhs 和 rhs 执行逐元素转换，生成 result 张量。

输入

输出

限制条件

• (C1) type(lhs) = type(rhs)。
• (C2) shape(result) = shape(lhs)。
• (C3) element_type(result) 的类型为 complex<E>，其中 E = element_type(lhs)。

示例

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs) : (tensor<2xf64>, tensor<2xf64>) -> tensor<2xcomplex<f64>>
// %result: [(1.0, 2.0), (3.0, 4.0)]

更多示例

复合型

语义

封装了由其他 StableHLO 操作组成的运算，接受 inputs 和 composite_attributes 并生成 results。相应操作的语义由 decomposition 属性实现。composite 操作可以替换为其分解，而不会更改程序语义。如果内嵌分解无法提供相同的操作语义，请优先使用 custom_call。

version 字段（默认为 0）用于表示复合语义何时发生变化。

输入

输出

限制条件

• (C1) is_namespaced_op_name(name)
• (C2) is_defined_in_parent_scope(decomposition)
• (C3) types(inputs...) == input_types(decomposition)
• (C4) types(results...) == output_types(decomposition)

示例

%results = "stablehlo.composite"(%input0, %input1) {
  name = "my_namespace.my_op",
  composite_attributes = {
    my_attribute = "my_value"
  },
  decomposition = @my_op,
  version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>

更多示例

concatenate

语义

沿 dimension 维度按给定实参的相同顺序串联 inputs，并生成 result 张量。更正式地说，result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1]，其中：

1. id = d0 + ... + dk-1 + kd。
2. d 等于 dimension，而 d0、... 是 inputs 的第 d 个维度的大小。

输入

输出

限制条件

• (C1) same(element_type(inputs...))。
• (C2) same(shape(inputs...))（dim(inputs..., dimension) 除外）。
• (C3) 0 < size(inputs)。
• (C4) 0 <= dimension < rank(inputs[0])。
• (C5) element_type(result) = element_type(inputs[0])。
• (C6) shape(result) = shape(inputs[0]) 但以下情况除外：
  • dim(result, dimension) = dim(inputs[0], dimension) + ...。

示例

// %input0: [[1, 2], [3, 4], [5, 6]]
// %input1: [[7, 8]]
%result = "stablehlo.concatenate"(%input0, %input1) {
  dimension = 0 : i64
} : (tensor<3x2xi64>, tensor<1x2xi64>) -> tensor<4x2xi64>
// %result: [[1, 2], [3, 4], [5, 6], [7, 8]]

更多示例

常量

语义

根据常量 value 生成 output 张量。

输入

输出

限制条件

• (C1) type(value) = type(output)。

示例

%output = "stablehlo.constant"() {
  value = dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>
} : () -> tensor<2x2xf32>
// %output: [[0.0, 1.0], [2.0, 3.0]]

更多示例

转化

语义

对 operand 张量执行从一种元素类型到另一种元素类型的逐元素转换，并生成 result 张量。

对于boolean-to-any-supported-type的转换，值 false 会转换为零，值 true 会转换为一。对于any-supported-type-to-boolean的转换，零值会转换为 false，非零值会转换为 true。如需了解此功能如何处理复杂类型，请参阅下文。

对于涉及整数到整数、整数到浮点数或浮点数到浮点数的转换，如果源值可以用目标类型精确表示，则结果值就是该精确表示。否则，行为尚待确定 (#180)。

对于涉及floating-point-to-integer的转换，系统会截断小数部分。如果截断后的值无法用目标类型表示，则行为待定 (#180)。

涉及复数到复数的转换遵循与浮点数到浮点数转换相同的行为，用于转换实部和虚部。

对于从复数到任何其他类型和从任何其他类型到复数的转换，系统会分别忽略源虚数值或将目标虚数值归零。实部的转换遵循浮点转换。

原则上，此操作可以表示反量化（从量化张量转换为常规张量）、量化（从常规张量转换为量化张量）和重新量化（在量化张量之间转换），但目前我们有专门的操作来处理这些情况 - uniform_dequantize 用于第一种使用情形，uniform_quantize 用于第二种和第三种使用情形。未来，这两个操作可能会合并为 convert (#1576)。

输入

输出

限制条件

• (C1) shape(operand) = shape(result)。

示例

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand) : (tensor<3xi64>) -> tensor<3xcomplex<f64>>
// %result: [(-1.0, 0.0), (0.0, 0.0), (1.0, 0.0)]

更多示例

卷积

语义

计算 lhs 的窗口与 rhs 的切片之间的点积，并生成 result。下图通过一个具体示例展示了如何根据 lhs 和 rhs 计算 result 中的元素。

更正式地说，请考虑从 lhs 的角度重新构建输入，以便能够表达 lhs 的窗口：

• lhs_window_dimensions = lhs_shape(dim(lhs, input_batch_dimension), dim(rhs, kernel_spatial_dimensions), dim(lhs, input_feature_dimension))。
• lhs_window_strides = lhs_shape(1, window_strides, 1)。
• lhs_padding = lhs_shape([0, 0], padding, [0, 0])。
• lhs_base_dilations = lhs_shape(1, lhs_dilation, 1)。
• lhs_window_dilations = lhs_shape(1, rhs_dilation, 1)。

此重构使用以下辅助函数：

• lhs_shape(n, hw, c) = permute([n] + hw + [c], [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension])。
• result_shape(n1, hw, c1) = permute([n1] + hw + [c1], [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension])。
• permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1]，其中 j[d] = i[permutation[d]]。

如果 feature_group_count = 1 和 batch_group_count = 1，则对于 index_space(dim(result, output_spatial_dimensions...)) 中的所有 output_spatial_index，result[result_shape(:, output_spatial_index, :)] = dot_product，其中：

• padding_value = constant(0, element_type(lhs))。
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1)。
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides。
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations)。
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true])。此功能似乎未被使用，因此我们计划在未来将其移除 (#1181)。
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension])。

如果 feature_group_count > 1：

• lhses = split(lhs, feature_group_count, input_feature_dimension)。
• rhses = split(rhs, feature_group_count, kernel_output_feature_dimension)。
• results... = convolution(lhses..., rhses..., ..., feature_group_count=1, ...)。
• result = concatenate(results, output_feature_dimension)。

如果 batch_group_count > 1：

• lhses = split(lhs, batch_group_count, input_batch_dimension)。
• rhses = split(rhs, batch_group_count, kernel_output_feature_dimension)。
• results... = convolution(lhses..., rhses..., ..., batch_group_count=1, ...)。
• result = concatenate(results, output_feature_dimension)。

对于量化类型，执行 dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result))。

对于混合量化类型，执行 hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs)。

输入

输出

限制条件

• (C1) N = rank(lhs) = rank(rhs)。
• (C2) size(window_strides) = N - 2。
• (C3) 0 < window_strides。
• (C4) shape(padding) = [N - 2, 2]。
• (C5) size(lhs_dilation) = N - 2。
• (C6) 0 < lhs_dilation。
• (C7) size(rhs_dilation) = N - 2。
• (C8) 0 < rhs_dilation。
• (C9) size(window_reversal) = N - 2。
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0。
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0。
• (C12) size(input_spatial_dimensions) = N - 2。
• (C13) 给定 input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]：
  • is_unique(input_dimensions)。
  • 0 <= input_dimensions < N。
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count。
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0。
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0。
• (C17) size(kernel_spatial_dimensions) = N - 2。
• (C18) 给定 kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]：
  • is_unique(kernel_dimensions)。
  • 0 <= kernel_dimensions < N。
• (C19) size(output_spatial_dimensions) = N - 2。
• (C20) 给定 output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]：
  • is_unique(output_dimensions)。
  • 0 <= output_dimensions < N。
• (C21) 0 < feature_group_count。
• (C22) 0 < batch_group_count。
• (C23) feature_group_count = 1 or batch_group_count = 1。
• (C24) size(precision_config) = 2。
• (C25) dim(result, result_dim) 的定义如下：
  • 如果 result_dim = output_batch_dimension，则为 dim(lhs, input_batch_dimension) / batch_group_count。
  • 如果 result_dim = output_feature_dimension，则为 dim(rhs, kernel_output_feature_dimension)。
  • 否则为 num_windows，其中：
  • output_spatial_dimensions[spatial_dim] = result_dim。
  • lhs_dim = input_spatial_dimensions[spatial_dim]。
  • rhs_dim = kernel_spatial_dimensions[spatial_dim]。
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1。
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1]。
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1。
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim]。
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1。
• (C26) rank(result) = N。
• 如果运算使用非量化张量：
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result)。
• 如果操作使用量化张量：
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)。
  • (C29) 如果 is_per_axis_quantized(rhs)，则 quantization_dimension(rhs) = kernel_output_feature_dimension。
  • (C30) 如果 is_per_axis_quantized(result)，则quantization_dimension(result) = output_feature_dimension。
  • 如果 is_quantized(lhs)：
  • (C31) storage_type(lhs) = storage_type(rhs)。
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)。
  • (C33) 如果 is_per_tensor_quantized(rhs)，则is_per_tensor_quantized(result)。
  • 如果 !is_quantized(lhs)：
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result)。

示例

// %lhs: [[
//        [
//          [1], [2], [5], [6]
//        ],
//        [
//          [3], [4], [7], [8]
//        ],
//        [
//          [10], [11], [14], [15]
//        ],
//        [
//          [12], [13], [16], [17]
//        ]
//      ]]
//
// %rhs: [
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]]
//       ]
%result = "stablehlo.convolution"(%lhs, %rhs) {
  window_strides = array<i64: 4, 4>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  // In the StableHLO dialect, dimension numbers are encoded via:
  // `[<input dimensions>]x[<kernel dimensions>]->[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" are spatial dimensions.
  dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  batch_group_count = 1 : i64,
  feature_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

更多示例

余弦

语义

对 operand 张量执行逐元素余弦运算，并生成 result 张量。根据元素类型，执行以下操作：

• 对于浮点数：来自 IEEE-754 的 cos。
• 对于复数：复余弦。
• 对于量化类型：dequantize_op_quantize(cosine, operand, type(result))。

输入

输出

限制条件

• (C1) baseline_type(operand) = baseline_type(result)。

示例

// %operand: [
//            [0.0, 1.57079632],       // [0, pi/2]
//            [3.14159265, 4.71238898] // [pi, 3pi/2]
//           ]
%result = "stablehlo.cosine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.0], [-1.0, 0.0]]

更多示例

count_leading_zeros

语义

逐元素计算 operand 张量中前导零位的数量，并生成 result 张量。

输入

输出

限制条件

• (C1) type(operand) = type(result)。

示例

// %operand: [[0, 1], [128, -1]]
%result = "stablehlo.count_leading_zeros"(%operand) : (tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[64, 63], [56, 0]]

更多示例

custom_call

语义

封装了实现定义的运算 call_target_name，该运算接受 inputs 和 called_computations 并生成 results。has_side_effect、backend_config 和 api_version 可用于提供其他实现定义的元数据。

目前，此操作包含相当杂乱的元数据集合，反映了 XLA 编译器中对应操作的自然演变。未来，我们计划统一此元数据 (#741)。

输入

输出

（XLA GPU 支持）特殊 custom_call 目标

有三个与 buffer 类型相关的特殊 call_target_name：CreateBuffer 创建未初始化的 buffer，Pin 创建已初始化的 buffer，Unpin 释放 buffer 并返回 buffer 的内容。

%uninitialized_buffer = "stablehlo.custom_call"() {
  call_target_name = "CreateBuffer",
  api_version = 4 : i32,
} : () -> memref<4xf64>

%initialized_buffer = "stablehlo.custom_call"(%init_value) {
  call_target_name = "Pin",
  api_version = 4 : i32,
} : (tensor<4xf64>) -> memref<4xf64>

%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
  call_target_name = "Unpin",
  api_version = 4 : i32,
} : (memref<4xf64>) -> tensor<4xf64>

别名

某些 custom_call 操作可能需要输出中的一部分和操作数中的一部分共享同一内存。这可以通过 output_operand_aliases 来表示。别名对表示法包含一个表示输出部分的输出元组索引列表，以及一个 operand_index 和一个表示运算数部分的运算数元组索引列表。如果相应类型不是 tuple 类型，则输出或操作数元组索引的列表为空；对于任意嵌套的元组类型，该列表可以任意长。这与 XLA 别名表示法类似。

别名对中的输出部分和输入部分必须具有相同的类型。对于不是对 CreateBuffer、Pin 和 Unpin 的调用的 custom_call 操作，buffer 实参最多可出现在一个别名对中，而 buffer 输出必须出现在一个别名对中。

示例

%results = "stablehlo.custom_call"(%input0) {
  call_target_name = "foo",
  has_side_effect = false,
  backend_config = {bar = 42 : i32},
  api_version = 4 : i32,
  called_computations = [@foo]
} : (tensor<f64>) -> tensor<f64>

%updated_buffer = "stablehlo.custom_call"(%buffer) {
  call_target_name = "Update",
  api_version = 4 : i32,
  output_operand_aliases = [
    #stablehlo.output_operand_alias<output_tuple_indices = [],
      operand_index = 0,
      operand_tuple_indices = []>]
} : (memref<4xf64>) -> memref<4xf64>

除

语义

对被除数 lhs 张量和除数 rhs 张量执行逐元素除法，并生成 result 张量。根据元素类型，执行以下操作：

• 对于整数：整数除法，生成代数商，并舍弃任何小数部分。
• 对于浮点数：来自 IEEE-754 的 division。
• 对于复数：复数除法。
• 对于量化类型：
  • dequantize_op_quantize(divide, lhs, rhs, type(result))。

输入

输出

限制条件

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)。

示例

// %lhs: [17.1, -17.1, 17.1, -17.1]
// %rhs: [3.0, 3.0, -3.0, -3.0]
%result = "stablehlo.divide"(%lhs, %rhs) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

更多示例

dot_general

语义

计算 lhs 的切片与 rhs 的切片之间的点积，并生成 result 张量。

更正式地说，result[result_index] = dot_product，其中：

• lhs_result_dimensions = [d for d in axes(lhs) and d not in lhs_batching_dimensions and d not in lhs_contracting_dimensions]。
• rhs_result_dimensions = [d for d in axes(rhs) and d not in rhs_batching_dimensions and d not in rhs_contracting_dimensions]。
• result_batching_index + result_lhs_index + result_rhs_index = result_index 其中 size(result_batching_index) = size(lhs_batching_dimensions)、size(result_lhs_index) = size(lhs_result_dimensions) 和 size(result_rhs_index) = size(rhs_result_dimensions)。
• transposed_lhs = transpose(lhs, lhs_batching_dimensions + lhs_result_dimensions + lhs_contracting_dimensions)。
• transposed_lhs_slice = slice(transposed_lhs, result_batching_index + result_lhs_index + [:, ..., :])。
• reshaped_lhs_slice = reshape(transposed_lhs_slice, dims(lhs, lhs_contracting_dimensions))。
• transposed_rhs = transpose(rhs, rhs_batching_dimensions + rhs_result_dimensions + rhs_contracting_dimensions)。
• transposed_rhs_slice = slice(transposed_rhs, result_batching_index + result_rhs_index + [:, ..., :])。
• reshaped_rhs_slice = reshape(transposed_rhs_slice, dims(rhs, rhs_contracting_dimensions))。
• dot_product = reduce( inputs=[multiply(reshaped_lhs_slice, reshaped_rhs_slice)], init_values=[constant(0, element_type(result))], dimensions=range(size(lhs_contracting_dimensions)), body=lambda x, y: add(x, y))。

对于量化类型，执行 dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result))。

对于混合量化类型，执行 hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs)。

precision_config 用于控制在加速器后端上进行计算时速度与准确性之间的权衡。可以是以下值之一（目前，这些枚举值的语义尚未完全指定，但我们计划在 #755 中解决此问题）：

• DEFAULT：计算速度最快，但对原始数字的近似程度最低。
• HIGH：计算速度较慢，但能更准确地逼近原始数字。
• HIGHEST：计算速度最慢，但最接近原始数字。

DotAlgorithm 用于定义实现点运算的算法的主要属性，同时还定义了精度。如果设置了算法属性字段，则 precision_config 必须为 DEFAULT。DotAlgorithms 没有默认值，因为默认参数是由实现定义的。因此，所有点算法字段都可以设置为 None，以指定空点算法，该算法将改用 precision_config 值。

DotAlgorithm 字段包括：

• lhs_precision_type 和 rhs_precision_type，运算的左侧和右侧四舍五入到的精度。精度类型与输入和输出的存储类型无关。
• accumulation_type 用于累积的精度。
• 当我们在执行将左侧实参和/或右侧实参分解为多个组成部分并对这些值执行多个“原始”点运算的算法时，会应用 lhs_component_count、rhs_component_count 和 num_primitive_operations - 通常是为了模拟更高的精度（例如利用 bfloat16 人工智能数据类型进行更高精度的计算：bf16_6x tf32_3x 等）。对于没有分解的算法，这些值应设置为 1。
• allow_imprecise_accumulation 用于指定是否允许在某些步骤（例如 CUBLASLT_MATMUL_DESC_FAST_ACCUM）中以较低精度进行累积。

DotAlgorithm 属性示例：

// Inputs are casted to tf32, and then accumulated in f32:
{lhs_precision_type = tf32,
 rhs_precision_type = tf32,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = false}


// bf16_6x: each input is decomposed to 3 bf16 components, then 6 dot operations are done on those components, and the result is accumulated in f32.
{lhs_precision_type = bf16,
 rhs_precision_type = bf16,
 accumulation_type = f32,
 lhs_component_count = 3,
 rhs_component_count = 3,
 num_primitive_operations = 6,
 allow_imprecise_accumulation = false}


// Inputs are (casted to) f8e5m2, and we accumulate in f32, but for some steps we may accumulate in lower precision.
{lhs_precision_type = f8e5m2,
 rhs_precision_type = f8e5m2,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = true}

具体支持哪些组合由实现方决定。一般来说，StableHLO 的使用者无法保证每种算法都受每种加速器类型的支持。如果不支持给定的算法，则应引发错误，而不是回退到替代算法。StableHLO 验证将尽力进行验证，防止使用已知在任何硬件上都不受支持的算法。

如需查看一些受支持的算法值，请参阅 xla_data.proto > Algorithm。工单 #2483 记录了创建关于后端支持的算法的集中式文档的计划。

输入

输出

限制条件

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions)。
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions)。
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions)。
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions)。
• (C5) 0 <= lhs_batching_dimensions < rank(lhs)。
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs)。
• (C7) 0 <= rhs_batching_dimensions < rank(rhs)。
• (C8) 0 <= rhs_contracting_dimensions < rank(rhs)。
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...)。
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...)。
• (C11) size(precision_config) = 2。
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions)。
• 如果运算使用非量化张量：
  • (C13) element_type(lhs) = element_type(rhs)。
• 如果操作使用量化张量：
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)。
  • (C15) zero_points(rhs) = 0。
  • (C16) 如果 is_per_axis_quantized(rhs)，则 quantization_dimension(rhs) 不在 rhs_contracting_dimensions 中。
  • 如果 is_quantized(lhs)：
  • (C17) storage_type(lhs) = storage_type(rhs)。
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)。
  • (C19) 如果 is_per_tensor_quantized(rhs)，则 is_per_tensor_quantized(result)。
  • 如果 !is_quantized(lhs)：
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result)。
• 如果 !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation)：
  • (C21) precision_config... = DEFAULT。
  • (C22) 0 < lhs_component_count。
  • (C23) 0 < rhs_component_count。
  • (C24) 0 < num_primitive_operations。

示例

// %lhs: [
//        [[1, 2],
//         [3, 4]],
//        [[5, 6],
//         [7, 8]]
//       ]
// %rhs: [
//        [[1, 0],
//         [0, 1]],
//        [[1, 0],
//         [0, 1]]
//       ]
%result = "stablehlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #stablehlo.dot<
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimensions = [1]
  >,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
  algorithm = #stablehlo.dot_algorithm<
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation = false
  >
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

更多示例

dynamic_broadcast_in_dim

语义

此操作在功能上与 broadcast_in_dim 操作相同，但结果形状是通过 output_dimensions 动态指定的。

该操作还接受可选属性 known_expanding_dimensions、known_nonexpanding_dimensions，以表达有关维度扩展行为的静态知识。如果未指定，则假定所有维度都可能扩大。

输入

输出

限制条件

• (C1) element_type(result) 由以下人员提供：
  • element_type(operand)（如果 !is_per_axis_quantized(operand)）。
  • element_type(operand)，但 quantization_dimension(operand)、scales(operand) 和 zero_points(operand) 可能分别与 quantization_dimension(result)、scales(result) 和 zero_points(result) 不同，否则。
• (C2) size(broadcast_dimensions) = rank(operand)。
• (C3) 0 <= broadcast_dimensions < rank(result)。
• (C4) is_unique(broadcast_dimensions)。
• (C5) 对于 axes(operand) 中的所有 d：
  • dim(operand, d) = 1 或
  • dim(operand, d) = dim(result, broadcast_dimensions[d])。
• (C6) 如果 is_per_axis_quantized(result)：
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)]。
  • 如果值为 dim(operand, quantization_dimension(operand)) = 1，则值为 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))。
• (C7) size(output_dimensions) = rank(result)。
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions)。
• (C9) 0 <= known_expanding_dimensions < rank(operand)。
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand)。

示例

// %operand: [
//            [1, 2, 3]
//           ]
%operand = stablehlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = stablehlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "stablehlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
  broadcast_dimensions = array<i64: 2, 1>,
  known_expanding_dimensions = array<i64: 0>,
  known_nonexpanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

更多示例

dynamic_conv

语义

此操作在功能上与卷积操作相同，但填充是通过 padding 动态指定的。

输入