StableHLO 是机器学习 (ML) 模型中高级操作 (HLO) 的一种操作集。StableHLO 充当不同机器学习框架和机器学习编译器之间的可移植层:生成 StableHLO 程序的机器学习框架与使用 StableHLO 程序的机器学习编译器兼容。
我们的目标是在各种机器学习框架(例如 TensorFlow、JAX 和 PyTorch)与机器学习编译器(例如 XLA 和 IREE)之间建立更高的互操作性,从而简化和加速机器学习开发。为此,本文档提供了 StableHLO 编程语言的规范。
本规范包含三个主要部分。首先,程序部分介绍了 StableHLO 程序的结构,这些程序由 StableHLO 函数组成,而这些函数本身由 StableHLO 操作组成。在该结构中,Ops 部分指定各个操作的语义。执行部分提供了所有这些操作在程序中一起执行的语义。最后,表示法部分讨论了整个规范中使用的表示法。
计划
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
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。
这定义了对尺寸大小是静态已知的静态形状的支持。将来,我们还计划支持动态形状,其中尺寸大小部分或完全未知 (#8)。此外,我们还计划探索将张量类型扩展到维度大小和元素类型之外,例如,包括布局 (#629) 和稀疏性 (#1078)。
QuantizedTensorType ::= 'tensor' '<' Shape QuantizedTensorElementType '>'
QuantizedTensorElementType ::= '!quant.uniform' '<'
QuantizationStorageType
['<' QuantizationStorageMin ':' QuantizationStorageMax '>']
':' QuantizationExpressedType
[':' QuantizationDimension]
',' QuantizationParameters '>'
QuantizationStorageType ::= IntegerType
QuantizationStorageMin ::= IntegerConstant
QuantizationStorageMax ::= IntegerConstant
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerConstant
QuantizationParameters ::= QuantizationParameter
| '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale ':' QuantizationZeroPoint
QuantizationScale ::= FloatConstant
QuantizationZeroPoint ::= IntegerConstant
名称 | 类型 | 限制条件 |
---|---|---|
storage_type |
整数类型 | (C1-C4)、(C9) |
storage_min |
整数常量 | (C2)、(C4)、(C8) |
storage_max |
整数常量 | (C3)、(C4)、(C8) |
expressed_type |
浮点类型 | (C1)、(C5) |
quantization_dimension |
可选整数常量 | (C11-C13) |
scales |
可变数的浮点数常量 | (C5-C7)、(C10)、(C11)、(C13) |
zero_points |
整数常量的不同数 | (C8-C10) |
量化元素类型表示介于 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)
num_bits(storage_type) < num_bits(expressed_type)
。 - (C2)
type(storage_min) = storage_type
。 - (C3)
type(storage_max) = storage_type
。 - (C4)
min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type)
。 - (C5)
type(scales...) = expressed_type
。 - (C6)
0 < scales
。 - (C7)
is_finite(scales...)
。 - (C8)
storage_min <= zero_points <= storage_max
。 - (C9)
type(zero_points...) = storage_type
。 - (C10)
size(scales) = size(zero_points)
。 - (C11) 如果
is_empty(quantization_dimension)
,则size(scales) = 1
。 - (C12)
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)
。
- (C12)
TokenType ::= 'token'
令牌类型表示令牌,即某些操作生成和使用的不透明值。令牌用于对操作强制执行执行顺序(如执行部分中所述)。
TupleType ::= 'tuple' '<' TupleElementTypes '>'
TupleElementTypes ::= [ValueType {',' ValueType}]
元组类型表示元组,即异构列表。元组是一项旧版功能,仅用于与 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 ::= 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f8E4M3FN' | 'f8E5M2' | 'f8E4M3FNUZ' | 'f8E5M2FNUZ'
| 'f8E4M3B11FNUZ' | 'bf16' | 'f16' | 'f32' | 'f64'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'
元素类型表示张量类型的元素。与许多编程语言不同,这些类型在 StableHLO 中不是顶级。这意味着,StableHLO 程序无法直接表示这些类型的值(因此,习惯使用 tensor<T>
类型的 0 维张量值表示 T
类型的标量值)。
- 布尔值类型表示布尔值
true
和false
。 - 整数类型可以是有符号 (
si
) 或无符号 (ui
),并且具有支持的位宽度(4
、8
、16
、32
或64
)之一。有符号siN
类型表示-2^(N-1)
到2^(N-1)-1
(含)之间的整数值,无符号uiN
类型表示从0
到2^N-1
(含)的整数值。 - 浮点类型可以是以下其中一项:
f8E4M3FN
和f8E5M2
类型,分别对应于适用于深度学习的 FP8 格式中所述的 FP8 格式的E4M3
和E5M2
编码。f8E4M3FNUZ
和f8E5M2FNUZ
类型,对应于深度神经网络的 8 位数值格式中所述的 FP8 格式的E4M3
和E5M2
编码。f8E4M3B11FNUZ
类型,对应于深度神经网络的混合式 8 位浮点 (HFP8) 训练和推断中所述的 FP8 格式的E4M3
编码。bf16
类型,对应于 BFloat16:Cloud TPU 上高性能的秘诀中所述的bfloat16
格式。f16
、f32
和f64
类型分别对应于 IEEE 754 标准中所述的binary16
(“半精度”)、binary32
(“单精度”)和binary64
(“双精度”)格式。
- 复杂类型表示同时具有同一元素类型的实部和虚部的复杂值。支持的复杂类型包括
complex<f32>
(两个部分均为f32
类型)和complex<f64>
(两个部分均为f64
类型)。
FunctionType ::= '(' InputTypes ')' '->' '(' OutputTypes ')'
InputTypes ::= [ValueType {',' ValueType}]
OutputTypes ::= [ValueType {',' ValueType}]
函数类型同时表示已命名函数和匿名函数。它们具有输入类型(->
左侧的类型列表)和输出类型(->
右侧的类型列表)。在许多编程语言中,函数类型是第一类,但未在 StableHLO 中。
StringType ::= 'string'
字符串类型表示字节序列。与许多编程语言不同,字符串类型不是 StableHLO 中的第一类,仅用于为程序元素指定静态元数据。
运维
“StableHLO 运算”(也称为“运算”)代表机器学习模型中一组封闭的高级运算。如上所述,StableHLO 语法深受 MLIR 启发,后者不一定是最符合人体工程学的替代方案,但可以说最适合 StableHLO 的目标是在机器学习框架和机器学习编译器之间提高互操作性。
Op ::= [OpOutputs] OpName OpInputs ':' OpSignature
OpName ::= '"' 'stablehlo' '.' OpMnemonic '"'
OpMnemonic ::= 'abs' | 'add' | ...
StableHLO 操作(也称为 ops)具有名称、输入/输出和签名。名称由 stablehlo.
前缀和一个助记符组成,用于唯一标识一项受支持的操作。如需查看所有受支持操作的完整列表,请参阅下文。
目前,实际的 StableHLO 程序有时包含本文档中未说明的操作。未来,我们计划将这些运算吸纳到 StableHLO 运算集中,或者禁止它们出现在 StableHLO 程序中。同时,下面列出了这些操作:
builtin.module
、func.func
、func.call
和func.return
(#425)。chlo
操作 (#602)。- StableHLO 操作中的“Not in HLO”类别 - 它们最初是 StableHLO 运算集的一部分,但后来被认为不适合它:
broadcast
、create_token
、cross-replica-sum
、dot
、einsum
、torch_index_select
、unary_einsum
(#3)。 - StableHLO 操作的“动态主义”类别 - 它们是从 MHLO 引导而来,但我们尚未确定它们:
compute_reshape_shape
、cstr_reshapable
、dynamic_broadcast_in_dim
、dynamic_conv
、dynamic_gather
、dynamic_iota
、dynamic_pad
、dynamic_reshape
、real_dynamic_slice
、set_dimension_size
(#8)。 - 形状计算,包括
arith
、shape
和tensor
操作 (#8)。
OpInputs ::= OpInputValues OpInputFuncs OpInputAttrs
OpInputValues ::= '(' [OpInputValue {',' OpInputValue}] ')'
OpInputValue ::= ValueId
OpInputFuncs ::= ['(' OpInputFunc {',' OpInputFunc} ')']
OpInputAttrs ::= ['{' OpInputAttr {',' OpInputAttr} '}']
OpOutputs ::= [OpOutput {',' OpOutput} '=']
OpOutput ::= ValueId
操作会使用输入并生成输出。输入分为输入值(在执行期间计算)、输入函数(以静态方式提供,因为在 StableHLO 函数中不是顶级值)和输入属性(也以静态方式提供)。操作使用和生成的输入和输出类型取决于其助记符。例如,add
操作使用 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 运算集,或禁止它们出现在 StableHLO 程序中。同时,下面列出了这些属性:
layout
(#629)。mhlo.frontend_attributes
(#628)。mhlo.sharding
(#619)。output_operand_aliases
(#740)。- 位置元数据 (#594)。
OpSignature ::= '(' [ValueType {',' ValueType}] ')' '->' '(' [ValueType {',' ValueType}] ')'
操作签名包含所有输入值的类型(->
左侧的类型列表)和所有输出值的类型(->
右侧的类型列表)。严格来说,输入类型是冗余的,而输出类型几乎总是冗余的(因为对于大多数 StableHLO 操作,输出类型可以从输入推断出来)。不过,为了与 MLIR 兼容,操作签名特意成为 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
。
Ops Agent 可以
abs
语义
对 operand
张量执行元素级抽象运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于有符号整数:整数模。
- 对于浮点数:IEEE-754 中的
abs
。 - 对于复数:复数模。
- 对于量化类型:
dequantize_op_quantize(abs, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
带符号整数、浮点或复杂类型的张量,或每个张量量化张量 | (C1-C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
带符号整数/浮点类型的张量或每个张量量化张量的张量 | (C1-C2) |
限制条件
- (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]
add
语义
执行两个张量 lhs
和 rhs
的元素级加法,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于布尔值:逻辑 OR。
- 对于整数:整数加法。
- 对于浮点数:IEEE-754 中的
addition
。 - 适用于复数:复数加法。
- 对于量化类型:
dequantize_op_quantize(add, lhs, rhs, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
或每张量量化张量 | (C1) |
(I2) | rhs |
或每张量量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1) |
限制条件
- (C1)
baseline_type(lhs) = baseline_type(rhs) = baseline_type(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
的数据依赖关系。
输入内容
标签 | 名称 | 类型 |
---|---|---|
(I1) | inputs |
token 的可变数 |
输出
名称 | 类型 |
---|---|
result |
token |
示例
// %input0: !stablehlo.token
// %input1: !stablehlo.token
%result = "stablehlo.after_all"(%input0, %input1) : (!stablehlo.token, !stablehlo.token) -> !stablehlo.token
all_gather
语义
在 StableHLO 进程网格中的每个进程组内,沿着 all_gather_dim
串联每个进程的 operand
张量值,并生成一个 result
张量。
该操作将 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
中:
- 针对
process_group
中的所有receiver
的operands@receiver = [operand@sender for sender in process_group]
权限。 - 针对
process_group
中的所有process
的result@process = concatenate(operands@process, all_gather_dim)
权限。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1)、(C6) |
(I2) | all_gather_dim |
si64 类型的常量 |
(C1)、(C6) |
(I3) | replica_groups |
si64 类型的二维张量常数 |
(C2-C4) |
(I4) | channel_id |
si64 类型的常量 |
(C5) |
(I5) | use_global_device_ids |
i1 类型的常量 |
(C5) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C6) |
限制条件
- (C1)
0 <= all_gather_dim < rank(operand)
。 - (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(result) = type(operand)
,但以下情况除外:dim(result, all_gather_dim) = dim(operand, all_gather_dim) * dim(process_groups, 1)
.
示例
// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
%result = "stablehlo.all_gather"(%operand) {
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<2x4xi64>
// %result@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
all_reduce
语义
在 StableHLO 进程网格中的每个进程组内,将归约函数 computation
应用于每个进程的 operand
张量值,并生成 result
张量。
该操作将 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
中:
result@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]))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C5)、(C6) |
(I2) | replica_groups |
si64 类型的一维张量常量的变数 |
(C1-C3) |
(I3) | channel_id |
si64 类型的常量 |
(C4) |
(I4) | use_global_device_ids |
i1 类型的常量 |
(C4) |
(I5) | computation |
function | (C5) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C6-C7) |
限制条件
- (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(result) = shape(operand)
。 - (C7)
element_type(result) = E
。
示例
// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [1, 2, 3, 4]
// %operand@(1, 0): [5, 6, 7, 8]
%result = "stablehlo.all_reduce"(%operand) ({
^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_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<i64>) -> tensor<i64>
// %result@(0, 0): [6, 8, 10, 12]
// %result@(1, 0): [6, 8, 10, 12]
all_to_all
语义
在 StableHLO 进程网格中的每个进程组内,将 operand
张量的值沿 split_dimension
拆分为多个部分,将拆分的部分分散到各个进程之间,沿着 concat_dimension
串联分散的部分并生成一个 result
张量。
该操作将 StableHLO 进程网格拆分为 process_groups
,其定义如下:
- 如果
channel_id <= 0
,则为cross_replica(replica_groups)
。 - 如果
channel_id > 0
,则为cross_partition(replica_groups)
。
之后,在每个 process_group
中:
- 针对
process_group
中的所有sender
使用split_parts@sender = split(operand@sender, split_count, split_dimension)
。 scattered_parts@receiver = [split_parts@sender[receiver_index] for sender in process_group]
,其中receiver_index = process_group.index(receiver)
。result@process = concatenate(scattered_parts@process, concat_dimension)
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1-C3)、(C9) |
(I2) | split_dimension |
si64 类型的常量 |
(C1)、(C2)、(C9) |
(I3) | concat_dimension |
si64 类型的常量 |
(C3)、(C9) |
(I4) | split_count |
si64 类型的常量 |
(C2)、(C4)、(C8)、(C9) |
(I5) | replica_groups |
si64 类型的二维张量常数 |
(C5-C8) |
(I6) | channel_id |
si64 类型的常量 |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C9) |
限制条件
- (C1)
0 <= split_dimension < rank(operand)
。 - (C2)
dim(operand, split_dimension) % split_count = 0
。 - (C3)
0 <= concat_dimension < rank(operand)
。 - (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(result) = type(operand)
,但以下情况除外:dim(result, split_dimension) = dim(operand, split_dimension) / split_count
.dim(result, concat_dimension) = dim(operand, concat_dimension) * split_count
.
示例
// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
// [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
// [13, 14, 15, 16]]
%result = "stablehlo.all_to_all"(%operand) {
split_dimension = 1 : i64,
concat_dimension = 0 : i64,
split_count = 2 : i64,
replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<2x4xi64>) -> tensor<4x2xi64>
// %result@(0, 0): [[1, 2],
// [5, 6],
// [9, 10],
// [13, 14]]
// %result@(1, 0): [[3, 4],
// [7, 8],
// [11, 12],
// [15, 16]]
和
语义
对两个张量 lhs
和 rhs
执行元素级 AND,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于布尔值:逻辑 AND。
- 对于整数:按位 AND。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
布尔值或整数类型的张量 | (C1) |
(I2) | rhs |
布尔值或整数类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
布尔值或整数类型的张量 | (C1) |
限制条件
- (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))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
(I2) | rhs |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (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
语义
计算从 grad_output
向外传播的 batch_norm_training
的多个输入的梯度,并生成 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))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点类型或每个张量量化张量的张量 | (C1-C3)、(C5) |
(I2) | scale |
浮点或每张量量化类型的一维张量 | (C2)、(C4)、(C5) |
(I3) | mean |
浮点或每张量量化类型的一维张量 | (C2)、(C4) |
(I4) | variance |
浮点或每张量量化类型的一维张量 | (C2)、(C4) |
(I5) | grad_output |
浮点类型或每个张量量化张量的张量 | (C2)、(C3) |
(I6) | epsilon |
f32 类型的常量 |
|
(I7) | feature_index |
si64 类型的常量 |
(C1)、(C5) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
grad_operand |
浮点类型或每个张量量化张量的张量 | (C2)、(C3) |
grad_scale |
浮点或每张量量化类型的一维张量 | (C2)、(C4) |
grad_offset |
浮点或每张量量化类型的一维张量 | (C2)、(C4) |
限制条件
- (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
语义
在除 feature_index
维度以外的所有维度上归一化 operand
张量,并生成 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))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点类型或每个张量量化张量的张量 | (C1-C7) |
(I2) | scale |
浮点或每张量量化类型的一维张量 | (C2)、(C3) |
(I3) | offset |
浮点或每张量量化类型的一维张量 | (C2)、(C4) |
(I4) | mean |
浮点或每张量量化类型的一维张量 | (C5) |
(I5) | variance |
浮点或每张量量化类型的一维张量 | (C2)、(C6) |
(I6) | epsilon |
f32 类型的常量 |
|
(I7) | feature_index |
si64 类型的常量 |
(C1)、(C3-C6) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点类型或每个张量量化张量的张量 | (C2)、(C7) |
限制条件
- (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))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点类型或每个张量量化张量的张量 | (C1) |
(I2) | scale |
浮点或每张量量化的一维张量 | (C2)、(C3) |
(I3) | offset |
浮点或每张量量化的一维张量 | (C2)、(C4) |
(I4) | epsilon |
f32 类型的常量 |
(C1)、(C3-C6) |
(I5) | feature_index |
si64 类型的常量 |
(C1)、(C3-C6) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
output |
浮点类型或每个张量量化张量的张量 | (C7) |
batch_mean |
浮点或每张量量化的一维张量 | (C2)、(C5) |
batch_var |
浮点或每张量量化的一维张量 | (C2)、(C6) |
限制条件
- (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
返回给定值的内存中表示法,其行为是由实现定义的,因为张量的精确表示法是由实现定义的,元素类型的确切表示法也是由实现定义的。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或量化张量 | (C1-C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或量化张量 | (C1-C2) |
限制条件
- (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
.- 针对所有
0 <= i < R
dim(result, i) = dim(operand, i)
。 dim(result, R) * num_bits(E') = num_bits(E)
.- 如果为
num_bits(E') > num_bits(E)
: rank(result) = R - 1
.- 针对所有
0 <= i < R
dim(result, i) = dim(operand, i)
。 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]]
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或量化张量 | (C1-C2)、(C5-C6) |
(I2) | broadcast_dimensions |
si64 类型的一维张量常数 |
(C2-C6) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或量化张量 | (C1)、(C3)、(C5-C6) |
限制条件
- (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]
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | index |
si32 类型的 0 维张量 |
|
(I2) | branches |
函数变量数 | (C1-C4) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变数量的张量、量化张量或词元 | (C4) |
限制条件
- (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))
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (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))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点类型或每个张量量化张量的张量 | (C1) |
限制条件
- (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
)矩阵的形式。对立三角形(即严格的上三角形或相应的严格下三角形)中的输出值是由实现定义的。
如果存在 i
,且输入矩阵不是埃尔米特正定矩阵,则此行为是未定义的。
对于量化类型,请执行 dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | a |
浮点或复杂类型或每个张量量化张量的张量 | (C1-C3) |
(I2) | lower |
i1 类型的 0 维张量常数 |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (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)。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | min |
或每张量量化张量 | (C1)、(C3) |
(I2) | operand |
或每张量量化张量 | (C1-C4) |
(I3) | max |
或每张量量化张量 | (C2)、(C3) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C4) |
限制条件
- (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(0, element_type(result)), [], type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
张量 | (C3) |
(I2) | replica_groups |
si64 类型的一维张量常量的变数 |
(C1)、(C2) |
(I3) | channel_id |
si64 类型的常量 |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
张量 | (C3) |
限制条件
- (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]
(如果存在使process_groups[i, 1] = process
如此的i
)。broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C5) |
(I2) | source_target_pairs |
si64 类型的二维张量常数 |
(C1-C4) |
(I3) | channel_id |
si64 类型的常量 |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1) |
限制条件
- (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]]
compare
语义
根据 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
运算。此功能似乎未使用,因此将来,我们计划将其移除 (#584)。
对于复杂元素类型,使用提供的 comparison_direction
和 compare_type
执行 (real, imag)
对的字典顺序比较。对复数施加排序涉及令人惊讶的语义,因此未来我们计划在 comparison_direction
为 GE
、GT
、LE
或 LT
时取消对复数的支持 (#560)。
适用于量化类型。执行 dequantize_compare(lhs, rhs,
comparison_direction)
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
或每张量量化张量 | (C1-C3) |
(I2) | rhs |
或每张量量化张量 | (C1-C2) |
(I3) | comparison_direction |
EQ 、NE 、GE 、GT 、LE 和 LT 的枚举 |
|
(I4) | compare_type |
FLOAT 、TOTALORDER 、SIGNED 和 UNSIGNED 的枚举 |
(C3) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
布尔值类型的张量 | (C2) |
限制条件
- (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
。 FLOAT
如果is_float(element_type(lhs))
,则选择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
张量。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
f32 或 f64 类型的张量 |
(C1-C3) |
(I2) | rhs |
f32 或 f64 类型的张量 |
(C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
复杂类型的张量 | (C2)、(C3) |
限制条件
- (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)]
concatenate
语义
按照与给定参数相同的顺序沿 dimension
维度串联 inputs
,并生成一个 result
张量。更正式地说,是 result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1]
,其中:
id = d0 + ... + dk-1 + kd
.d
等于dimension
,而d0
, ... 是第d
个维度大小 (inputs
)。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | inputs |
可变张量数量或每个张量量化张量 | (C1-C6) |
(I2) | dimension |
si64 类型的常量 |
(C2)、(C4)、(C6) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C5-C6) |
限制条件
- (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
张量。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | value |
常量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
output |
或量化张量 | (C1) |
限制条件
- (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
会转换为 1。对于any-supported-type-to-boolean的转换,零值转换为 false
,非零值转换为 true
。请参阅下文,了解对于复杂类型是如何工作的。
对于涉及“整数到整数”“整数到浮点数”或“浮点到浮点数”的转换,如果源值可以在目标类型中精确表示,则结果值就是确切的表示法。否则,行为将待定 (#180)。
对于涉及floating-point-to-integer的转换,小数部分会被截断。如果截断的值无法在目的地类型中表示,则该行为待定 (#180)。
涉及复数到复数的转换遵循与浮点到浮点转换相同的行为,用于转换实部和虚部。
对于“复数到任何其他类型”和“任何其他类型到复数”complex-to-any-other-type转化,系统会分别忽略来源虚值或目标虚值值为零。complex-to-any-other-type实际部分的转换在浮点数转换之后。
原则上,此操作可以表示反量化(从量化张量转换为常规张量)、量化(从常规张量转换为量化张量)和重新量化(量化张量之间的转换),但目前我们有专门的操作 - uniform_dequantize
用于第一个用例,uniform_quantize
用于第二个和第三个用例。将来,这两个操作可能会合并到 convert
中 (#1576)。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
张量 | (C1) |
限制条件
- (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))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
或每张量量化张量 | (C1)、(C10-C11)、(C14) (C25)、(C27-C30) |
(I2) | rhs |
或量化张量 | (C1)、(C14-C16)、(C25)、(C27-C32) |
(I3) | window_strides |
si64 类型的一维张量常数 |
(C2-C3)、(C25) |
(I4) | padding |
si64 类型的二维张量常数 |
(C4)、(C25) |
(I5) | lhs_dilation |
si64 类型的一维张量常数 |
(C5-C6)、(C25) |
(I6) | rhs_dilation |
si64 类型的一维张量常数 |
(C7-C8)、(C25) |
(I7) | window_reversal |
i1 类型的一维张量常数 |
(C9) |
(I8) | input_batch_dimension |
si64 类型的常量 |
(C10)、(C13)、(C25) |
(I9) | input_feature_dimension |
si64 类型的常量 |
(C11)、(C13-C14) |
(I10) | input_spatial_dimensions |
si64 类型的一维张量常数 |
(C12)、(C13)、(C25) |
(I11) | kernel_input_feature_dimension |
si64 类型的常量 |
(C14)、(C18) |
(I12) | kernel_output_feature_dimension |
si64 类型的常量 |
(C15-C16)、(C18)、(C25)、(C32) |
(I13) | kernel_spatial_dimensions |
si64 类型的一维张量常数 |
(C17-C18)、(C25) |
(I14) | output_batch_dimension |
si64 类型的常量 |
(C20)、(C25) |
(I15) | output_feature_dimension |
si64 类型的常量 |
(C20)、(C25)、(C33) |
(I16) | output_spatial_dimensions |
si64 类型的一维张量常数 |
(C19-C20)、(C25) |
(I17) | feature_group_count |
si64 类型的常量 |
(C11)、(C14)、(C16)、(C21)、(C23) |
(I18) | batch_group_count |
si64 类型的常量 |
(C10)、(C15)、(C22)、(C23)、(C25) |
(I19) | precision_config |
DEFAULT 、HIGH 和 HIGHEST 的枚举数量不等 |
(C24) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或量化张量 | (C25-C28)、(C30-C31)、(C33) |
限制条件
- (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)
。
- (C27)
- 如果操作使用量化张量:
- (C28)
is_quantized_tensor(lhs) and is_quantized_tensor(rhs) and is_quantized_tensor(result)
。 - (C29)
storage_type(lhs) = storage_type(rhs)
。 - (C30)
expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)
。 - (C31) 如果为
is_per_tensor_quantized(rhs)
,则设为is_per_tensor_quantized(result)
。 - (C32) 如果为
is_per_axis_quantized(rhs)
,则设为quantization_dimension(rhs) = kernel_output_feature_dimension
。 - (C33) 如果为
is_per_axis_quantized(result)
,则设为quantization_dimension(result) = output_feature_dimension
。
- (C28)
示例
// %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 = dense<4> : tensor<2xi64>,
padding = dense<0> : tensor<2x2xi64>,
lhs_dilation = dense<2> : tensor<2xi64>,
rhs_dilation = dense<1> : tensor<2xi64>,
window_reversal = dense<false> : tensor<2xi1>,
// 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]>,
feature_group_count = 1 : i64,
batch_group_count = 1 : i64,
precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi32>
// %result: [[
// [[10], [26]],
// [[46], [62]]
// ]]
余弦
语义
对 operand
张量执行元素级余弦运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
cos
。 - 对于复数:复余弦。
- 对于量化类型:
dequantize_op_quantize(cosine, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (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
张量。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
整数类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数类型的张量 | (C1) |
限制条件
- (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)。
输入内容
标签 | 名称 | 类型 |
---|---|---|
(I1) | inputs |
值数量不等 |
(I2) | call_target_name |
string 类型的常量 |
(I3) | has_side_effect |
i1 类型的常量 |
(I4) | backend_config |
string 类型的常量 |
(I5) | api_version |
si32 类型的常量 |
(I6) | called_computations |
string 类型的常量可变数 |
输出
名称 | 类型 |
---|---|
results |
值数量不等 |
示例
%results = "stablehlo.custom_call"(%input0) {
call_target_name = "foo",
has_side_effect = false,
backend_config = "bar",
api_version = 1 : i32,
called_computations = [@foo]
} : (tensor<f64>) -> tensor<f64>
除号
语义
对被除数 lhs
和除数 rhs
张量进行元素级除法,并生成 result
张量。根据元素类型,执行以下操作:
- 对于整数:整数除法,可生成代数商,并舍弃所有小数部分。
- 对于浮点数:IEEE-754 中的
division
。 - 适用于复数:复数除法。
- 对于量化类型:
dequantize_op_quantize(divide, lhs, rhs, type(result))
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
(I2) | rhs |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
限制条件
- (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))
。
这仅指定了每个张量量化的语义。正在进行每轴量化 (#1574)。此外,将来,我们可能会考虑添加对混合量化的支持 (#1575)。
precision_config
用于控制加速器后端上的计算在速度和准确性之间的权衡。这可能是以下某种原因(目前,这些枚举值的语义尚未明确,但我们计划在 #755 中解决此问题):
DEFAULT
:以最快的速度计算,但与原始数字最接近的近似值最低。HIGH
:计算速度更慢,但更接近原始数字。HIGHEST
:计算最慢,但与原始数字最接近的近似值。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
或每张量量化张量 | (C5-C6)、(C9-C10)、(C12-C16) |
(I2) | rhs |
或每张量量化张量 | (C7-C10)、(C12) |
(I3) | lhs_batching_dimensions |
si64 类型的一维张量常数 |
(C1)、(C3)、(C5)、(C9)、(C12) |
(I4) | rhs_batching_dimensions |
si64 类型的一维张量常数 |
(C1)、(C4)、(C7)、(C9) |
(I5) | lhs_contracting_dimensions |
si64 类型的一维张量常数 |
(C2)、(C3)、(C6)、(C10) |
(I6) | rhs_contracting_dimensions |
si64 类型的一维张量常数 |
(C2)、(C4)、(C8)、(C10) |
(I7) | precision_config |
DEFAULT 、HIGH 和 HIGHEST 的枚举数量不等 |
(C11) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C12)、(C14)、(C16) |
限制条件
- (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)
。
- (C13)
- 如果操作使用量化张量:
- (C14)
is_quantized(lhs) and is_quantized(rhs) and is_quantized(result)
。 - (C15)
storage_type(lhs) = storage_type(rhs)
。 - (C16)
expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)
。 - (C17)
zero_points(rhs) = 0
。
- (C14)
示例
// %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>]
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
// [[1, 2],
// [3, 4]],
// [[5, 6],
// [7, 8]]
// ]
dynamic_slice
语义
使用动态计算的起始索引从 operand
中提取 Slice,并生成 result
张量。start_indices
包含每个可能调整的维度的切片起始索引,slice_sizes
包含每个维度的切片大小。更正式地说,是 result[result_index] = operand[operand_index]
,其中:
adjusted_start_indices = clamp(0, start_indices, shape(operand) - slice_sizes)
.operand_index = adjusted_start_indices + result_index
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1)、(C2)、(C4) |
(I2) | start_indices |
整数类型的零维张量的可变数 | (C2)、(C3) |
(I3) | slice_sizes |
si64 类型的一维张量常数 |
(C2)、(C4)、(C5) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1)、(C5) |
限制条件
- (C1)
element_type(operand) = element_type(result)
。 - (C2)
size(start_indices) = size(slice_sizes) = rank(operand)
。 - (C3)
same(type(start_indices...))
。 - (C4)
0 <= slice_sizes <= shape(operand)
。 - (C5)
shape(result) = slice_sizes
。
示例
// %operand: [
// [0, 0, 1, 1],
// [0, 0, 1, 1],
// [0, 0, 0, 0],
// [0, 0, 0, 0]
// ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_slice"(%operand, %start_indices0, %start_indices1) {
slice_sizes = dense<[2, 2]> : tensor<2xi64>
} : (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>
// %result: [
// [1, 1],
// [1, 1]
// ]
dynamic_update_slice
语义
生成一个 result
张量,它等于 operand
张量,只不过使用 update
中的值更新以 start_indices
开始的 Slice。更正式地说,result[result_index]
的定义如下:
- 如果
0 <= update_index < shape(update)
,则为update[update_index]
,其中:adjusted_start_indices = clamp(0, start_indices, shape(operand) - shape(update))
.update_index = result_index - adjusted_start_indices
.
- 否则为
operand[result_index]
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1-C4)、(C6) |
(I2) | update |
或每张量量化张量 | (C2)、(C3)、(C6) |
(I3) | start_indices |
整数类型的零维张量的可变数 | (C4)、(C5) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1) |
限制条件
- (C1)
type(operand) = type(result)
。 - (C2)
element_type(update) = element_type(operand)
。 - (C3)
rank(update) = rank(operand)
。 - (C4)
size(start_indices) = rank(operand)
。 - (C5)
same(type(start_indices...))
。 - (C6)
0 <= shape(update) <= shape(operand)
。
示例
// %operand: [
// [1, 1, 0, 0],
// [1, 1, 0, 0],
// [1, 1, 1, 1],
// [1, 1, 1, 1]
// ]
// %update: [
// [1, 1],
// [1, 1]
// ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_update_slice"(%operand, %update, %start_indices0, %start_indices1)
: (tensor<4x4xi32>, tensor<2x2xi32>, tensor<i64>, tensor<i64>) -> tensor<4x4xi32>
// %result: [
// [1, 1, 1, 1],
// [1, 1, 1, 1],
// [1, 1, 1, 1],
// [1, 1, 1, 1]
// ]
指数函数
语义
对 operand
张量执行元素级指数运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
exp
。 - 适用于复数:复数指数。
- 对于量化类型:
dequantize_op_quantize(exponential, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.exponential"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[1.0, 2.7182818284590451], [7.3890560989306504, 20.085536923187668]]
exponential_minus_one
语义
对 operand
张量执行元素级指数减一运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
expm1
。 - 对于复数:复数指数减 1。
- 对于量化类型:
dequantize_op_quantize(exponential_minus_one, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [0.0, 1.0]
%result = "stablehlo.exponential_minus_one"(%operand) : (tensor<2xf64>) -> tensor<2xf64>
// %result: [0.0, 1.71828187]
FF
语义
对真实和复杂输入/输出执行正向和逆逆傅里叶转换。
fft_type
是以下值之一:
FFT
:转发复杂到复杂 FFT。IFFT
:复杂到复杂 FFT 的反函数。RFFT
:转发实际到复杂 FFT。IRFFT
:实数转换为复数 FFT(即取复数,返回实数)。
更正式地说,假设函数 fft
将复杂类型的一维张量作为输入,生成相同类型的一维张量作为输出,并计算离散傅里叶转换:
对于 fft_type = FFT
,result
定义为一系列 L 计算的最终结果,其中 L = size(fft_length)
。例如,对于 L = 3
:
result1[i0, ..., :] = fft(operand[i0, ..., :])
.result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1])
.result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1])
.
此外,假设函数 ifft
具有相同的类型签名,并计算 fft
的逆:
对于 fft_type = IFFT
,result
定义为 fft_type = FFT
的反转。例如,对于 L = 3
:
result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1])
.result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1])
.result[i0, ..., :] = ifft(result2[i0, ..., :])
.
此外,假设函数 rfft
接受浮点类型的一维张量,会生成相同浮点语义的复杂类型的一维张量,其工作原理如下:
rfft(real_operand) = truncated_result
,其中complex_operand... = (real_operand..., 0.0)
.complex_result = fft(complex_operand)
.truncated_result = complex_result[:(rank(complex_result) / 2 + 1)]
.
(在针对实际运算数计算离散傅里叶转换时,结果的前 N/2 + 1
元素明确定义结果的其余部分,因此 rfft
的结果会被截断以避免计算冗余元素)。
对于 fft_type = RFFT
,result
定义为一系列 L 计算的最终结果,其中 L = size(fft_length)
。例如,对于 L = 3
:
result1[i0, ..., :] = rfft(operand[i0, ..., :])
.result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1])
.result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1])
.
最后,假设函数 irfft
具有相同的类型签名,并计算 rfft
的逆:
对于 fft_type = IRFFT
,result
定义为 fft_type = RFFT
的反转。例如,对于 L = 3
:
result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1])
.result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1])
.result[i0, ..., :] = irfft(result2[i0, ..., :])
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型的张量 | (C1)、(C2)、(C4)、(C5) |
(I2) | fft_type |
FFT 、IFFT 、RFFT 和 IRFFT 的枚举 |
(C2)、(C5) |
(I3) | fft_length |
si64 类型的一维张量常数 |
(C1)、(C3)、(C4) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型的张量 | (C2)、(C4)、(C5) |
限制条件
- (C1)
size(fft_length) <= rank(operand)
。 - (C2)
operand
和result
元素类型之间的关系有所不同:- 如果为
fft_type = FFT
,则element_type(operand)
和element_type(result)
具有相同的复杂类型。 - 如果为
fft_type = IFFT
,则element_type(operand)
和element_type(result)
具有相同的复杂类型。 - 如果为
fft_type = RFFT
,则element_type(operand)
是浮点类型,而element_type(result)
是相同浮点语义的复杂类型。 - 如果为
fft_type = IRFFT
,则element_type(operand)
是复杂类型,而element_type(result)
是相同浮点语义的浮点类型。
- 如果为
- (C3)
1 <= size(fft_length) <= 3
。 - (C4) 如果在
operand
和result
之间有一个浮点类型的张量real
,则shape(real)[-size(fft_length):] = fft_length
。 - (C5)
shape(result) = shape(operand)
,但以下几项除外:- 如果为
fft_type = RFFT
,则返回dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1
。 - 如果为
fft_type = IRFFT
,则返回dim(operand, -1) = dim(result, -1) = 0 ? 0 : dim(result, -1) / 2 + 1
。
- 如果为
示例
// %operand: [(1.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
%result = "stablehlo.fft"(%operand) {
fft_type = #stablehlo<fft_type FFT>,
fft_length = dense<4> : tensor<1xi64>
} : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
// %result: [(1.0, 0.0), (1.0, 0.0), (1.0, 0.0), (1.0, 0.0)]
floor
语义
执行 operand
张量的元素级平面,并生成一个 result
张量。
实现符合 IEEE-754 规范的 roundToIntegralTowardNegative
操作。对于量化类型,请执行 dequantize_op_quantize(floor, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.floor"(%operand) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-1.0, -1.0, 0.0, 0.0, 2.0]
收集
语义
从 start_indices
中指定的偏移量收集 operand
张量的切片,并生成一个 result
张量。
通过一个具体示例,下图显示了 result
中的元素如何映射到 operand
中的元素。该图选择了几个示例 result
索引,并详细介绍了它们对应的 operand
索引。
更正式地说,是 result[result_index] = operand[operand_index]
,其中:
batch_dims = [d for d in axes(result) and d not in offset_dims]
.batch_index = result_index[batch_dims...]
.start_index
的定义如下:start_indices[bi0, ..., :, ..., biN]
,其中bi
是batch_index
中的各个元素,如果index_vector_dim
<rank(start_indices)
,则会在index_vector_dim
索引处插入:
。- 否则为
[start_indices[batch_index]]
。
- 对于
axes(operand)
中的d_operand
,- 如果
d_operand = start_index_map[d_start]
,则为full_start_index[d_operand] = clamp(start_index[d_start], 0, dim(operand, d_operand) - slice_sizes[d_operand])
。 - 否则为
full_start_index[d_operand] = 0
。
- 如果
offset_index = result_index[offset_dims...]
.full_offset_index = [oi0, ..., 0, ..., oiN]
,其中oi
是offset_index
中的各个元素,0
是在collapsed_slice_dims
中的索引处插入。operand_index = full_start_index + full_offset_index
。
如果 indices_are_sorted
为 true
,则实现可以假定 start_indices
相对于 start_index_map
进行排序,否则行为将处于未定义状态。更正式地说,适用于 indices(result)
中的所有 i1 < i2
和 full_start_index(i1) <= full_start_index(i2)
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1)、(C7)、(C10-C12)、(C14) |
(I2) | start_indices |
整数类型的张量 | (C2)、(C3)、(C13) |
(I3) | offset_dims |
si64 类型的一维张量常数 |
(C1)、(C4-C5)、(C13) |
(I4) | collapsed_slice_dims |
si64 类型的一维张量常数 |
(C1)、(C6-C8)、(C13) |
(I5) | start_index_map |
si64 类型的一维张量常数 |
(C3)、(C9)、(C10) |
(I6) | index_vector_dim |
si64 类型的常量 |
(C2)、(C3)、(C13) |
(I7) | slice_sizes |
si64 类型的一维张量常数 |
(C8)、(C11-C13) |
(I8) | indices_are_sorted |
i1 类型的常量 |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C5)、(C13-C14) |
限制条件
- (C1)
rank(operand) = size(offset_dims) + size(collapsed_slice_dims)
。 - (C2)
0 <= index_vector_dim <= rank(start_indices)
。 - (C3)
size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1
。 - (C4)
is_unique(offset_dims) and is_sorted(offset_dims)
。 - (C5)
0 <= offset_dims < rank(result)
。 - (C6)
is_unique(collapsed_slice_dims) and is_sorted(collapsed_slice_dims)
。 - (C7)
0 <= collapsed_slice_dims < rank(operand)
。 - (C8)
slice_sizes[collapsed_slice_dims...] <= 1
。 - (C9)
is_unique(start_index_map)
。 - (C10)
0 <= start_index_map < rank(operand)
。 - (C11)
size(slice_sizes) = rank(operand)
。 - (C12)
0 <= slice_sizes <= shape(operand)
。 - (C13)
shape(result) = combine(batch_dim_sizes, offset_dim_sizes)
,其中:batch_dim_sizes = shape(start_indices)
,只不过不包含与index_vector_dim
对应的start_indices
的尺寸大小。offset_dim_sizes = shape(slice_sizes)
,区别在于不包含slice_sizes
中与collapsed_slice_dims
对应的尺寸尺寸。combine
将batch_dim_sizes
放置在与batch_dims
对应的轴上,并将offset_dim_sizes
放置在与offset_dims
对应的轴上。
- (C14)
element_type(operand) = element_type(result)
。
示例
// %operand: [
// [[1, 2], [3, 4], [5, 6], [7, 8]],
// [[9, 10],[11, 12], [13, 14], [15, 16]],
// [[17, 18], [19, 20], [21, 22], [23, 24]]
// ]
// %start_indices: [
// [[0, 0], [1, 0], [2, 1]],
// [[0, 1], [1, 1], [0, 2]]
// ]
%result = "stablehlo.gather"(%operand, %start_indices) {
dimension_numbers = #stablehlo.gather<
offset_dims = [2, 3],
collapsed_slice_dims = [0],
start_index_map = [1, 0],
index_vector_dim = 2>,
slice_sizes = dense<[1, 2, 2]> : tensor<3xi64>,
indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
// %result: [
// [
// [[1, 2], [3, 4]],
// [[3, 4], [5, 6]],
// [[13, 14], [15, 16]]
// ],
// [
// [[9, 10], [11, 12]],
// [[11, 12], [13, 14]],
// [[17, 18], [19, 20]]
// ]
// ]
get_dimension_size
语义
生成 operand
的指定 dimension
的大小。更正式地说,是 result = dim(operand, dimension)
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
张量 | (C1) |
(I2) | dimension |
si64 类型的常量 |
(C1) |
输出
名称 | 类型 |
---|---|
result |
si32 类型的 0 维张量 |
限制条件
- (C1)
0 <= dimension < rank(operand)
。
示例
// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.get_dimension_size"(%operand) {
dimension = 1 : i64
} : (tensor<2x3xi64>) -> tensor<i32>
// %result: 3
get_tuple_element
语义
提取 operand
元组的 index
位置处的元素,并生成一个 result
。更正式地说,result = operand[index]
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
tuple | (C1)、(C2) |
(I2) | index |
si32 类型的常量 |
(C1)、(C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
任何受支持的类型 | (C2) |
限制条件
- (C1)
0 <= index < size(operand)
。 - (C2)
type(result) = tuple_element_types(operand)[index]
。
示例
// %operand: ([1.0, 2.0], (3))
%result = "stablehlo.get_tuple_element"(%operand) {
index = 0 : i32
} : (tuple<tensor<2xf32>, tuple<tensor<i32>>>) -> tensor<2xf32>
// %result: [1.0, 2.0]
if
语义
根据 pred
的值,通过正好执行 true_branch
或 false_branch
中的一个函数来生成输出。更正式地说,result =
pred ? true_branch() : false_branch()
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | pred |
i1 类型的 0 维张量 |
|
(I2) | true_branch |
function | (C1-C3) |
(I3) | false_branch |
function | (C1)、(C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变数量的张量、量化张量或词元 | (C3) |
限制条件
- (C1)
input_types(true_branch) = input_types(false_branch) = []
。 - (C2)
output_types(true_branch) = output_types(false_branch)
。 - (C3)
type(results...) = output_types(true_branch)
。
示例
// %result_true_branch: 10
// %result_false_branch: 11
// %pred: true
%result = "stablehlo.if"(%pred) ({
"stablehlo.return"(%result_true_branch) : (tensor<i32>) -> ()
}, {
"stablehlo.return"(%result_false_branch) : (tensor<i32>) -> ()
}) : (tensor<i1>) -> tensor<i32>
// %result: 10
Imag
语义
从 operand
按元素提取虚部,并生成一个 result
张量。更正式地说,对于每个 x
元素:imag(x) = is_complex(x) ? imaginary_part(x) :
constant(0, element_type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型的张量 | (C1)、(C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点类型的张量 | (C1)、(C2) |
限制条件
- (C1)
shape(result) = shape(operand)
。 - (C2)
element_type(result)
的定义如下:- 如果
is_complex(operand)
,则为complex_element_type(element_type(operand))
。 - 否则为
element_type(operand)
。
- 如果
示例
// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.imag"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [2.0, 4.0]
信息流
语义
从信息流读取数据并生成 results
。
infeed_config
的语义是由实现定义的。
results
由载荷值在前和令牌值组成。将来,我们计划将载荷和令牌拆分为两个单独的输出以提高清晰度 (#670)。
输入内容
标签 | 名称 | 类型 |
---|---|---|
(I1) | token |
token |
(I2) | infeed_config |
string 类型的常量 |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变数量的张量、量化张量或词元 | (C1-C3) |
限制条件
- (C1)
0 < size(results)
。 - (C2)
is_empty(result[:-1])
或is_tensor(type(results[:-1]))
。 - (C3)
is_token(type(results[-1]))
。
示例
// %token: !stablehlo.token
// infeed_queue[0]: [[1, 2], [3, 4]]
// infeed_queue[1]: [[5, 6], [7, 8]]
%results0:2 = "stablehlo.infeed"(%token) {
infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results0#0: [[1, 2], [3, 4]]
%results1:2 = "stablehlo.infeed"(%token) {
infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results1#0: [[5, 6], [7, 8]]
Iota
语义
使用 iota_dimension
维度上从零开始递增的值填充 output
张量。更正式地说,
output[result_index] = constant(is_quantized(output) ?
quantize(result_index[iota_dimension], element_type(output)) :
result_index[iota_dimension], element_type(output))
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | iota_dimension |
si64 |
(C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
output |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
限制条件
- (C1)
0 <= iota_dimension < rank(output)
。
示例
%output = "stablehlo.iota"() {
iota_dimension = 0 : i64
} : () -> tensor<4x5xi32>
// %output: [
// [0, 0, 0, 0, 0],
// [1, 1, 1, 1, 1],
// [2, 2, 2, 2, 2],
// [3, 3, 3, 3, 3]
// ]
%output = "stablehlo.iota"() {
iota_dimension = 1 : i64
} : () -> tensor<4x5xi32>
// %output: [
// [0, 1, 2, 3, 4],
// [0, 1, 2, 3, 4],
// [0, 1, 2, 3, 4],
// [0, 1, 2, 3, 4]
// ]
is_finite
语义
执行元素级检查 x
中的值是否为有限值(即既不是 +Inf、-Inf 也不是 NaN),并生成 y
张量。实现符合 IEEE-754 规范的 isFinite
操作。对于量化类型,结果始终为 true
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | x |
浮点类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
y |
布尔值类型的张量 | (C1) |
限制条件
- (C1)
shape(x) = shape(y)
。
示例
// Logical values: -Inf, +Inf, NaN, ...
// %x: [0xFFF0000000000000, 0x7FF0000000000000, 0x7FF8000000000000, -10.0, -0.0, 0.0, 10.0]
%y = "stablehlo.is_finite"(%x) : (tensor<7xf64) -> tensor<7xi1>
// %y: [false, false, false, true, true, true, true]
日志
语义
对 operand
张量执行元素级对数运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
log
。 - 对于复数:复数对数。
- 对于量化类型:
dequantize_op_quantize(log, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [[1.0, 2.0], [3.0, 4.0]]
%result = "stablehlo.log"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.0, 0.69314718055994529], [1.0986122886681098, 1.3862943611198906]]
log_plus_one
语义
对 operand
张量执行元素级对数加一次运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
logp1
。 - 适用于复数:复数对数加 1。
- 对于量化类型:
dequantize_op_quantize(log_plus_one, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [0.0, -0.999, 7.0, 6.38905621, 15.0]
%result = "stablehlo.log_plus_one"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [0.0, -6.90776825, 2.07944155, 2.0, 2.77258873]
逻辑
语义
对 operand
张量执行元素级逻辑运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
division(1, addition(1, exp(-x)))
。 - 适用于复数:复数逻辑。
- 对于量化类型:
dequantize_op_quantize(logistic, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.logistic"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.5, 0.73105858], [0.88079708, 0.95257413]]
map
语义
沿着 dimensions
将映射函数 computation
应用于 inputs
,并生成一个 result
张量。
更正式地说,是 result[result_index] = computation(inputs...[result_index])
。
请注意,dimensions
目前未使用,日后可能会被移除 (#487)。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | inputs |
可变张量数量或每个张量量化张量 | (C1-C4) |
(I2) | dimensions |
si64 类型的一维张量常数 |
(C3) |
(I3) | computation |
function | (C4) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1)、(C4) |
限制条件
- (C1)
shape(inputs...) = shape(result)
。 - (C2)
0 < size(inputs) = N
。 - (C3)
dimensions = range(rank(inputs[0]))
。 - (C4)
computation
的类型为(tensor<E0>, ..., tensor<EN-1>) -> tensor<E'>
,其中Ei = element_type(inputs[i])
和E' = element_type(result)
。
示例
// %input0: [[0, 1], [2, 3]]
// %input1: [[4, 5], [6, 7]]
%result = "stablehlo.map"(%input0, %input1) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = stablehlo.multiply %arg0, %arg1 : tensor<i64>
stablehlo.return %0 : tensor<i64>
}) {
dimensions = dense<[0, 1]> : tensor<2xi64>
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[0, 5], [12, 21]]
最大值
语义
对 lhs
和 rhs
张量执行元素级最大运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于布尔值:逻辑 OR。
- 对于整数:最大值为整数。
- 对于浮点数:IEEE-754 中的
maximum
。 - 对于复数:
(real, imaginary)
对的字典顺序最大值。对复数施加排序涉及令人惊讶的语义,因此未来我们计划取消对此操作的支持 (#560)。 - 对于量化类型:
dequantize_op_quantize(maximum, lhs, rhs, type(result))
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
或每张量量化张量 | (C1) |
(I2) | rhs |
或每张量量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1) |
限制条件
- (C1)
baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)
。
示例
// %lhs: [[1, 2], [7, 8]]
// %rhs: [[5, 6], [3, 4]]
%result = "stablehlo.maximum"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 6], [7, 8]]
最小值
语义
对 lhs
和 rhs
张量执行元素级最小运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于布尔值:逻辑 AND。
- 对于整数:最小值为整数。
- 对于浮点数:IEEE-754 中的
minimum
。 - 对于复数:
(real, imaginary)
对的字典顺序最小值。对复数施加排序涉及令人惊讶的语义,因此未来我们计划取消对此操作的支持 (#560)。 - 对于量化类型:
dequantize_op_quantize(minimum, lhs, rhs, type(result))
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
或每张量量化张量 | (C1) |
(I2) | rhs |
或每张量量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1) |
限制条件
- (C1)
baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)
。
示例
// %lhs: [[1, 2], [7, 8]]
// %rhs: [[5, 6], [3, 4]]
%result = "stablehlo.minimum"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[1, 2], [3, 4]]
乘号
语义
执行两个张量 lhs
和 rhs
的元素级积,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于布尔值:逻辑 AND。
- 对于整数:整数乘法。
- 对于浮点数:IEEE-754 中的
multiplication
。 - 适用于复数:复数乘法。
- 对于量化类型:
dequantize_op_quantize(multiply, lhs, rhs, type(result))
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
或每张量量化张量 | (C1) |
(I2) | rhs |
或每张量量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.multiply"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 12], [21, 32]]
negate
语义
对 operand
张量执行元素级否定运算并生成 result
张量。根据元素类型,执行以下操作:
- 对于有符号整数:整数否定。
- 对于无符号整数:位转换为有符号整数、整数否定、位转换回无符号整数。
- 对于浮点数:IEEE-754 中的
negate
。 - 对于复数:复数否定。
- 对于量化类型:
dequantize_op_quantize(negate, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// Negation operation with integer Tensors
// %operand: [0, -2]
%result = "stablehlo.negate"(%operand) : (tensor<2xi32>) -> tensor<2xi32>
// %result: [0, 2]
// Negation operation with with complex tensors
// %operand: (2.5, 0.0)
%result = "stablehlo.negate"(%operand) : (tensor<1xcomplex<f32>>) -> tensor<1xcomplex<f32>>
// %result: [-2.5, -0.0]
非
语义
对 operand
张量执行非元素级操作,并生成一个 result
张量。
根据元素类型,执行以下操作:
- 对于布尔值:逻辑 NOT。
- 对于整数:按位 NOT。
参数
名称 | 类型 | 限制条件 |
---|---|---|
operand |
布尔值或整数类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
布尔值或整数类型的张量 | (C1) |
限制条件
- (C1)
type(operand) = type(result)
。
示例
// Bitwise operation with with integer tensors
// %operand: [[1, 2], [3, 4]]
%result = "stablehlo.not"(%operand) : (tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[-2, -3], [-4, -5]]
// Bitwise operation with with boolean tensors
// %operand: [true, false]
%result = "stablehlo.not"(%operand) : (tensor<2xi1>) -> tensor<2xi1>
// %result: [false, true]
optimization_barrier
语义
确保生成 operand
的操作先于任何依赖于 result
的操作执行,并防止编译器转换跨屏障移动操作。除此之外,操作是身份,即 result = operand
。
参数
名称 | 类型 | 限制条件 |
---|---|---|
operand |
可变张量数量、每个张量量化张量或词元 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
可变张量数量、每个张量量化张量或词元 | (C1) |
限制条件
- (C1)
type(operand...) = type(result...)
。
示例
// %operand0: 0.0
// %operand1: 1.0
%result0, %result1 = "stablehlo.optimization_barrier"(%operand0, %operand1) : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
// %result0: 0.0
// %result1: 1.0
或
语义
对两个张量 lhs
和 rhs
执行元素级 OR 运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于布尔值:逻辑 OR。
- 对于整数:按位 OR。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
整数或布尔值类型的张量 | (C1) |
(I2) | rhs |
整数或布尔值类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数或布尔值类型的张量 | (C1) |
限制条件
- (C1)
type(lhs) = type(rhs) = type(result)
。
示例
// Bitwise operation with with integer tensors
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 6], [7, 12]]
// Logical operation with with boolean tensors
// %lhs: [[false, false], [true, true]]
// %rhs: [[false, true], [false, true]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi1>, tensor<2x2xi1>) -> tensor<2x2xi1>
// %result: [[false, true], [true, true]]
外播
语义
将 inputs
写入输出并生成 result
词元。
outfeed_config
的语义是由实现定义的。
输入内容
标签 | 名称 | 类型 |
---|---|---|
(I1) | inputs |
可变数量的张量或量化张量 |
(I2) | token |
token |
(I3) | outfeed_config |
string 类型的常量 |
输出
名称 | 类型 |
---|---|
result |
token |
示例
%result = "stablehlo.outfeed"(%inputs0, %token) {
outfeed_config = ""
} : (tensor<2x2x2xi64>, !stablehlo.token) -> !stablehlo.token
垫
语义
通过在张量周围以及具有给定 padding_value
的张量元素之间填充内边距来扩展 operand
。
edge_padding_low
和 edge_padding_high
分别指定在每个维度的低端(索引 0 旁边)和高端(最高索引旁边)处添加的内边距。内边距可以为负数,其中负内边距的绝对值表示要从指定维度中移除的元素数量。
interior_padding
用于指定每个维度的任意两个元素之间添加的内边距大小(不得为负数)。内部内边距发生在边缘内边距之前,这样一来,负边缘内边距将会移除内内边距运算数中的元素。
更正式地说,result[result_index]
的定义如下:
- 如果
result_index = edge_padding_low + operand_index * (interior_padding + 1)
,则为operand[operand_index]
。 - 否则为
padding_value
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1)、(C2)、(C4) |
(I2) | padding_value |
0 维张量或每张量量化张量 | (C1) |
(I3) | edge_padding_low |
si64 类型的一维张量常数 |
(C1)、(C4) |
(I4) | edge_padding_high |
si64 类型的一维张量常数 |
(C1)、(C4) |
(I5) | interior_padding |
si64 类型的一维张量常数 |
(C2-C4) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C3-C6) |
限制条件
- (C1)
element_type(operand) = element_type(padding_value) = element_type(result)
。 - (C2)
size(edge_padding_low) = size(edge_padding_high) = size(interior_padding) = rank(operand)
。 - (C3)
0 <= interior_padding
。 - (C4)
shape(result) = shape(operand) + edge_padding_low + max(shape(operand) - 1, 0) * interior_padding + edge_padding_high
。
示例
// %operand: [
// [1, 2, 3],
// [4, 5, 6]
// ]
// %padding_value: 0
%result = "stablehlo.pad"(%operand, %padding_value) {
edge_padding_low = dense<[0, 1]> : tensor<2xi64>,
edge_padding_high = dense<[2, 1]> : tensor<2xi64>,
interior_padding = dense<[1, 2]> : tensor<2xi64>
} : (tensor<2x3xi32>, tensor<i32>) -> tensor<5x9xi32>
// %result: [
// [0, 1, 0, 0, 2, 0, 0, 3, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0],
// [0, 4, 0, 0, 5, 0, 0, 6, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0],
// [0, 0, 0, 0, 0, 0, 0, 0, 0]
// ]
partition_id
语义
生成当前进程的 partition_id
。
输出
名称 | 类型 |
---|---|
result |
ui32 类型的 0 维张量 |
示例
%result = "stablehlo.partition_id"() : () -> tensor<ui32>
爆破
语义
对 operand
张量中设置的位数执行元素级计数,并生成一个 result
张量。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
整数类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数类型的张量 | (C1) |
限制条件
- (C1)
type(operand) = type(result)
。
示例
// %operand: [0, 1, 2, 127]
%result = "stablehlo.popcnt"(%operand) : (tensor<4xi64>) -> tensor<4xi64>
// %result: [0, 1, 1, 7]
幂数
语义
通过 rhs
张量对 lhs
张量执行元素级指数运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于整数:整数指数。
- 对于浮点数:IEEE-754 中的
pow
。 - 对于复数:复数指数。
- 对于量化类型:
dequantize_op_quantize(power, lhs, rhs, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
(I2) | rhs |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %lhs: [-2.0, -0.0, -36.0, 5.0, 3.0, 10000.0]
// %rhs: [2.0, 2.0, 1.1, 2.0, -1.0, 10.0]
%result = "stablehlo.power"(%lhs, %rhs) : (tensor<6xf64>, tensor<6xf64>) -> tensor<6xf64>
// %result: [4.0, 0.0, -nan, 25.0, 0.333333343, inf]
real
语义
从 operand
按元素提取实际部分,并生成一个 result
张量。更正式地说,对于每个 x
元素:real(x) = is_complex(x) ? real_part(x) : x
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型的张量 | (C1)、(C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点类型的张量 | (C1)、(C2) |
限制条件
- (C1)
shape(result) = shape(operand)
。 - (C2)
element_type(result)
的定义如下:- 如果
is_complex(operand)
,则为complex_element_type(element_type(operand))
。 - 否则为
element_type(operand)
。
- 如果
示例
// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.real"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [1.0, 3.0]
接收
语义
使用 channel_id
接收来自频道的数据并生成 results
。
如果 is_host_transfer
为 true
,则操作会从主机传输数据。否则,它会从其他设备传输数据。这意味着,是由实现定义的。该标志与 channel_type
中提供的信息重复,因此将来我们计划只保留其中一个 (#666)。
results
由载荷值在前和令牌值组成。将来,我们计划将载荷和令牌拆分为两个单独的输出以提高清晰度 (#670)。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | token |
token |
(C4) |
(I2) | channel_id |
si64 类型的常量 |
|
(I3) | channel_type |
DEVICE_TO_DEVICE 和 HOST_TO_DEVICE 的枚举 |
(C1) |
(I4) | is_host_transfer |
i1 类型的常量 |
(C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变数量的张量、量化张量或词元 | (C2-C4) |
限制条件
- (C1)
channel_type
的定义如下:- 如果
is_host_transfer = true
,则为HOST_TO_DEVICE
; - 否则为
DEVICE_TO_DEVICE
。
- 如果
- (C2)
0 < size(results)
。 - (C3)
is_empty(result[:-1])
或is_tensor(type(results[:-1]))
。 - (C4)
is_token(type(results[-1]))
。
示例
%results0, %results1 = "stablehlo.recv"(%token) {
channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>,
is_host_transfer = true
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
reduce
语义
沿着 dimensions
对 inputs
和 init_values
应用归约函数 body
,并生成 results
张量。
归约顺序由实现定义,这意味着 body
和 init_values
必须形成单声道,以确保相应运算针对所有实现上的所有输入生成相同的结果。但是,这个条件不适用于许多常见的缩减。例如,body
的浮点加法和 init_values
的 0 实际上不会形成单声道,因为浮点加法不遵守结合律。
更正式地说,是 results...[j0, ..., jR-1] = reduce(input_slices_converted)
,其中:
input_slices = inputs...[j0, ..., :, ..., jR-1]
,其中:
插入dimensions
处。input_slices_converted = to_destination_type(input_slices..., type(func_inputs(body)[:len(func_inputs(body))//2])...)
.init_values_converted = to_destination_type(init_values..., type(func_inputs(body)[len(func_inputs(body))//2:])...)
.reduce(input_slices_converted) = exec(schedule)
,适用于某些二元树schedule
,其中:exec(node) = body(exec(node.left), exec(node.right))
.exec(leaf) = leaf.value
.
schedule
是实现定义的完整二元树,其有序遍历包括:input_slices_converted...[index]
值,适用于index_space(input_slices_converted)
中的所有index
(按index
的字典顺序升序)。- 在实现定义的位置插入一个实现定义数量的
init_values_converted
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | inputs |
可变张量数量或每个张量量化张量 | (C1-C4)、(C6)、(C7) |
(I2) | init_values |
零维张量或每张量量化张量的可变数 | (C2)、(C3) |
(I3) | dimensions |
si64 类型的一维张量常数 |
(C4)、(C5)、(C7) |
(I4) | body |
function | (C6) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变张量数量或每个张量量化张量 | (C3)、(C7)、(C8) |
限制条件
- (C1)
same(shape(inputs...))
。 - (C2)
element_type(inputs...) = element_type(init_values...)
。 - (C3)
0 < size(inputs) = size(init_values) = size(results) = N
。 - (C4)
0 <= dimensions < rank(inputs[0])
。 - (C5)
is_unique(dimensions)
。 - (C6)
body
的类型为(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ...,
tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)
,其中is_promotable(element_type(inputs[i]), Ei)
。 - (C7)
shape(results...) = shape(inputs...)
,只不过不包含与dimensions
对应的inputs...
的尺寸大小。 - (C8) 针对
[0,N)
中的所有i
使用element_type(results[i]) = Ei
。
示例
// %input = [[0, 1, 2, 3, 4, 5]]
// %init_value = 0
%result = "stablehlo.reduce"(%input, %init_value) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
dimensions = dense<1> : tensor<1xi64>
} : (tensor<1x6xi64>, tensor<i64>) -> tensor<1xi64>
// %result = [15]
reduce_precision
语义
执行 operand
到另一个使用 exponent_bits
和 mantissa_bits
的浮点类型的元素级转换,然后再转换回原始浮点类型并生成 output
张量。
更正式一点:
- 原始值的尾数位会更新,以将原始值四舍五入为可使用
mantissa_bits
使用roundToIntegralTiesToEven
语义表示的最接近的值。 - 然后,如果
mantissa_bits
小于原始值的尾数位,则尾数位会被截断为mantissa_bits
。 - 然后,如果中间结果的指数位不在
exponent_bits
提供的范围内,则中间结果会使用原始符号溢出到无穷大,或使用原始符号下溢到零。 - 对于量化类型,执行
dequantize_op_quantize( lambda operand: reduce_precision(operand, exponent_bits, mantissa_bits), operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点类型或每个张量量化张量的张量 | (C1) |
(I2) | exponent_bits |
si32 类型的常量 |
(C2) |
(I3) | mantissa_bits |
si32 类型的常量 |
(C3) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
output |
浮点类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(output)
。 - (C2)
1 <= exponent_bits
。 - (C3)
0 <= mantissa_bits
。
示例
// Logical values: +Inf, NaN, +Denormal, 0.0, 65519.0, 65520.0
// %operand: [0x7FF0000000000000, 0x7FFFFFFFFFFFFFFF, 0x0000000000000001, 0.0, 65519.0, 65520.0]
%output = "stablehlo.reduce_precision"(%operand) {
exponent_bits = 5 : i32,
mantissa_bits = 10 : i32
} : (tensor<6xf64>) -> tensor<6xf64>
// Logical values: +Inf, NaN, 0.0, 0.0, 65504.0, +Inf
// %output: [0x7FF0000000000000, 0x7FFFFFFFFFFFFFFF, 0.0, 0.0, 65504.0, 0x7FF0000000000000]
reduce_scatter
语义
在 StableHLO 进程网格中的每个进程组内,使用 computations
对每个进程中的 operand
张量值执行归约,沿着 scatter_dimension
将归约结果拆分为部分,并在进程之间散布拆分的部分,以生成 result
。
该操作将 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
中:
reduced_value = all_reduce(operand, replica_groups, channel_id, use_global_device_ids, computation)
.parts@sender = split(reduced_value@sender, dim(process_groups, 1), scatter_dimension)
.result@receiver = parts@sender[receiver_index]
- 针对process_group
中的所有sender
,其中receiver_index = process_group.index(receiver)
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1)、(C2)、(C7)、(C8) |
(I2) | scatter_dimension |
si64 类型的常量 |
(C1)、(C2)、(C8) |
(I3) | replica_groups |
si64 类型的二维张量常数 |
(C3-C5) |
(I4) | channel_id |
si64 类型的常量 |
(C6) |
(I5) | use_global_device_ids |
i1 类型的常量 |
(C6) |
(I6) | computation |
function | (C7) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C8-C9) |
限制条件
- (C1)
dim(operand, scatter_dimension) % dim(process_groups, 1) = 0
。 - (C2)
0 <= scatter_dimension < rank(operand)
。 - (C3)
is_unique(replica_groups)
。 - (C4)
size(replica_groups)
的定义如下:- 如果使用
cross_replica
,则为num_replicas
。 - 如果使用
cross_replica_and_partition
,则为num_replicas
。 - 如果使用
flattened_ids
,则为num_processes
。
- 如果使用
- (C5)
0 <= replica_groups < size(replica_groups)
。 - (C6) 如果
use_global_device_ids = true
,则channel_id > 0
。 - (C7)
computation
的类型为(tensor<E>, tensor<E>) -> (tensor<E>)
,其中is_promotable(element_type(operand), E)
。 - (C8)
shape(result) = shape(operand)
,但以下情况除外:dim(result, scatter_dimension) = dim(operand, scatter_dimension) / dim(process_groups, 1)
.
- (C9)
element_type(result) = E
。
示例
// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
// [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
// [13, 14, 15, 16]]
%result = "stablehlo.reduce_scatter"(%operand) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
scatter_dimension = 1 : i64,
replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<2x4xi64>) -> tensor<2x2xi64>
//
// %result@(0, 0): [[10, 12],
// [18, 20]]
// %result@(1, 0): [[14, 16],
// [22, 24]]
reduce_window
语义
将归约函数 body
应用于 inputs
和 init_values
的窗口,并生成 results
。
下图显示了如何使用一个具体示例根据 inputs...
计算 results...
中的元素。
更正式地说,是 results...[result_index] = reduce(windows, init_values, axes(inputs...), body)
(请参阅减少),其中:
padded_inputs = pad(inputs..., init_values..., padding[:, 0], padding[:, 1], base_dilations - 1)
.window_start = result_index * window_strides
.window_end = window_start + (window_dimensions - 1) * window_dilations + 1
.windows = slice(padded_inputs..., window_start, window_end, window_dilations)
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | inputs |
可变张量数量或每个张量量化张量 | (C1-C4)、(C6)、(C8)、(C10)、(C12)、(C13)、(C15) |
(I2) | init_values |
零维张量或每张量量化张量的可变数 | (C1)、(C13) |
(I3) | window_dimensions |
si64 类型的一维张量常数 |
(C4)、(C5)、(C15) |
(I4) | window_strides |
si64 类型的一维张量常数 |
(C6)、(C7)、(C15) |
(I5) | base_dilations |
si64 类型的一维张量常数 |
(C8)、(C9)、(C15) |
(I6) | window_dilations |
si64 类型的一维张量常数 |
(C10)、(C11)、(C15) |
(I7) | padding |
si64 类型的二维张量常数 |
(C12)、(C15) |
(I8) | body |
function | (C13) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变张量数量或每个张量量化张量 | (C1)、(C14-C16) |
限制条件
- (C1)
0 < size(inputs) = size(init_values) = size(results) = N
。 - (C2)
same(shape(inputs...))
。 - (C3)
element_type(inputs...) = element_type(init_values...)
。 - (C4)
size(window_dimensions) = rank(inputs[0])
。 - (C5)
0 < window_dimensions
。 - (C6)
size(window_strides) = rank(inputs[0])
。 - (C7)
0 < window_strides
。 - (C8)
size(base_dilations) = rank(inputs[0])
。 - (C9)
0 < base_dilations
。 - (C10)
size(window_dilations) = rank(inputs[0])
。 - (C11)
0 < window_dilations
。 - (C12)
shape(padding) = [rank(inputs[0]), 2]
。 - (C13)
body
的类型为(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ...,
tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)
,其中is_promotable(element_type(inputs[i]), Ei)
。 - (C14)
same(shape(results...))
。 - (C15)
shape(results[0]) = num_windows
,其中:dilated_input_shape = shape(inputs[0]) = 0 ? 0 : (shape(inputs[0]) - 1) * base_dilations + 1
.padded_input_shape = padding[:, 0] + dilated_input_shape + padding[:, 1]
.dilated_window_shape = (window_dimensions - 1) * window_dilations + 1
.is_empty_window = padded_input_shape = 0 || dilated_window_shape > padded_input_shape
.num_windows = is_empty_window ? 0 : floor((padded_input_shape - dilated_window_shape) / window_strides) + 1
.
- (C16) 针对
[0,N)
中的所有i
使用element_type(results[i]) = Ei
。
示例
// %input = [[1, 2], [3, 4], [5, 6]]
// %init_value = 0
%result = "stablehlo.reduce_window"(%input, %init_value) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
window_dimensions = dense<[2, 1]> : tensor<2xi64>,
window_strides = dense<[4, 1]> : tensor<2xi64>,
base_dilations = dense<[2, 1]> : tensor<2xi64>,
window_dilations = dense<[3, 1]> : tensor<2xi64>,
padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi64>, tensor<i64>) -> tensor<2x2xi64>
// %result = [[0, 0], [3, 4]]
余数
语义
执行被除数 lhs
和除数 rhs
张量的元素级余数,并生成 result
张量。
更正式地说,结果的符号取自被除数,结果的绝对值始终小于除数的绝对值。余数按 lhs - d * rhs
计算,其中 d
的计算公式如下:
- 对于整数:
stablehlo.divide(lhs, rhs)
。 - 对于浮点数:基于 IEEE-754 的
division(lhs, rhs)
,具有舍入属性roundTowardZero
。 - 对于复数:待定 (#997)。
- 对于量化类型:
dequantize_op_quantize(remainder, lhs, rhs, type(result))
.
对于浮点元素类型,此运算与 IEEE-754 规范中的 remainder
运算相反,其中 d
是最接近 lhs/rhs
的精确值且等于偶数的整数值。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
(I2) | rhs |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %lhs: [17, -17, 17, -17]
// %rhs: [3, 3, -3, -3]
%result = "stablehlo.remainder"(%lhs, %rhs) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi64>
// %result: [2, -2, 2, -2]
replica_id
语义
生成当前进程的 replica_id
。
输出
名称 | 类型 |
---|---|
result |
ui32 类型的 0 维张量 |
示例
%result = "stablehlo.replica_id"() : () -> tensor<ui32>
调整形状
语义
将 operand
张量重塑为 result
张量。从概念上讲,这相当于保留相同的规范化表示法,但可能会改变形状,例如从 tensor<2x3xf32>
更改为 tensor<3x2xf32>
或 tensor<6xf32>
。
更正式地说,是 result[result_index] = operand[operand_index]
,其中 result_index
和 operand_index
在 index_space(result)
和 index_space(operand)
的字典顺序中具有相同的位置。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或量化张量 | (C1-C3) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或量化张量 | (C1-C3) |
限制条件
- (C1)
element_type(result)
的计算公式如下:element_type(operand)
(如果!is_per_axis_quantized(operand)
)。element_type(operand)
,但quantization_dimension(operand)
和quantization_dimension(result)
可能会有所不同。
- (C2)
size(operand) = size(result)
。 - (C3) 如果
is_per_axis_quantized(operand)
:reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y)
.dim(operand, quantization_dimension(operand)) = dim(result, quantization_dimension(result))
.reduce(dims(operand, [quantization_dimension(operand) + 1, ..., rank(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [quantization_dimension(result) + 1, ..., rank(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y)
.
示例
// %operand: [[1, 2, 3], [4, 5, 6]]]
%result = "stablehlo.reshape"(%operand) : (tensor<2x3xi32>) -> tensor<3x2xi32>
// %result: [[1, 2], [3, 4], [5, 6]]
reverse
语义
逆转 operand
中元素的顺序沿指定的 dimensions
方向,并生成一个 result
张量。更正式地说,是 result[result_index] = operand[operand_index]
,其中:
operand_index[d] = dim(result, d) - result_index[d] - 1
(如果dimensions
为d
)。- 否则为
operand_index[d] = result_index[d]
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1)、(C3) |
(I2) | dimensions |
si64 类型的一维张量常数 |
(C2)、(C3) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1)、(C3) |
限制条件
- (C1)
type(operand) = type(result)
。 - (C2)
is_unique(dimensions)
。 - (C3)
0 <= dimensions < rank(result)
。
示例
// %operand = [[1, 2], [3, 4], [5, 6]]
%result = "stablehlo.reverse"(%operand) {
dimensions = dense<1> : tensor<1xi64>
} : (tensor<3x2xi32>) -> tensor<3x2xi32>
// %result: [[2, 1], [4, 3], [6, 5]]
广播
语义
使用 rng_distribution
算法生成随机数字,并生成给定形状 shape
的 result
张量。
如果为 rng_distribution = UNIFORM
,则按照间隔 [a, b)
内的均匀分布生成随机数字。如果为 a >= b
,则行为未定义。
如果为 rng_distribution = NORMAL
,则按照正态分布生成随机数字,其中平均值 = a
,标准差 = b
。如果为 b < 0
,则行为未定义。
生成随机数的确切方式由实现定义。例如,它们不一定具有确定性,不一定使用隐藏状态。
在与许多利益相关方的对话中,我们发现此操作已被有效地废弃,因此未来我们计划探索如何将其移除 (#597)。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | a |
整数、布尔值或浮点类型的 0 维张量 | (C1)、(C2) |
(I2) | b |
整数、布尔值或浮点类型的 0 维张量 | (C1)、(C2) |
(I3) | shape |
si64 类型的一维张量常数 |
(C3) |
(I4) | rng_distribution |
UNIFORM 和 NORMAL 的枚举 |
(C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数、布尔值或浮点类型的张量 | (C1-C3) |
限制条件
- (C1)
element_type(a) = element_type(b) = element_type(result)
。 - (C2) 如果
rng_distribution = NORMAL
,则is_float(a)
。 - (C3)
shape(result) = shape
。
示例
// %a = 0
// %b = 2
// %shape = [3, 3]
%result = "stablehlo.rng"(%a, %b, %shape) {
rng_distribution = #stablehlo<rng_distribution UNIFORM>
} : (tensor<i32>, tensor<i32>, tensor<2xi64>) -> tensor<3x3xi32>
// %result: [
// [1, 0, 1],
// [1, 1, 1],
// [0, 0, 0]
// ]
rng_bit_generator
语义
在给定初始状态 initial_state
的情况下,使用伪随机数生成器算法 rng_algorithm
返回填充了统一随机位的 output
和更新后的输出状态 output_state
。输出可以保证是 initial_state
的确定性函数,但在实现之间不一定是确定的。
rng_algorithm
是以下值之一:
DEFAULT
:实现定义的算法。THREE_FRY
:实现定义的 Threefry 算法变体。*PHILOX
:实现定义的 Philox 算法变体。*
* 请参阅:Salmon 等人,SC 2011。并行随机数字:简单到 1、2、3。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | rng_algorithm |
DEFAULT 、THREE_FRY 和 PHILOX 的枚举 |
(C2) |
(I2) | initial_state |
ui64 类型的一维张量 |
(C1)、(C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
output_state |
ui64 类型的一维张量 |
(C1) |
output |
整数或浮点类型的张量 |
限制条件
- (C1)
type(initial_state) = type(output_state)
。 - (C2)
size(initial_state)
的定义如下:- 如果
rng_algorithm = DEFAULT
,则为实现定义。 - 如果
rng_algorithm = THREE_FRY
,则为2
。 2
如果rng_algorithm = PHILOX
,则选择3
。
- 如果
示例
// %initial_state: [1, 2]
%output_state, %output = "stablehlo.rng_bit_generator"(%initial_state) {
rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
} : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x2xui64>)
// %output_state: [1, 6]
// %output: [
// [9236835810183407956, 16087790271692313299],
// [18212823393184779219, 2658481902456610144]
// ]
round_nearest_afz
语义
在 operand
张量上对最接近的整数执行元素级舍入,使关系远离 0,并生成 result
张量。实现符合 IEEE-754 规范的 roundToIntegralTiesToAway
操作。对于量化类型,请执行 dequantize_op_quantize(round_nearest_afz, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_afz"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-3.0, 0.0, 1.0, 1.0, 3.0]
round_nearest_even
语义
在 operand
张量上对最接近的整数执行元素级舍入,断开与偶数整数的关系,并生成 result
张量。实现符合 IEEE-754 规范的 roundToIntegralTiesToEven
操作。对于量化类型,请执行 dequantize_op_quantize(round_nearest_even, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_even"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-2.0, 0.0, 0.0, 1.0, 2.0]
rsqrt
语义
对 operand
张量执行元素级倒数平方根运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
rSqrt
。 - 对于复数:复数倒数平方根。
- 对于量化类型:
dequantize_op_quantize(rsqrt, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [[1.0, 4.0], [9.0, 25.0]]
%result = "stablehlo.rsqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.5], [0.33333343, 0.2]]
scatter
语义
生成 results
张量,这些张量等于 inputs
张量,但 scatter_indices
指定的多个切片是使用 update_computation
的值 updates
进行更新。
通过一个具体示例,下图显示了 updates...
中的元素如何映射到 results...
中的元素。该图选择了几个示例 updates...
索引,并详细介绍了它们对应的 results...
索引。
更正式地说,对于 index_space(updates[0])
中的所有 update_index
:
update_scatter_dims = [d for d in axes(updates[0]) and d not in update_window_dims]
.update_scatter_index = update_index[update_scatter_dims...]
.start_index
的定义如下:scatter_indices[si0, ..., :, ..., siN]
,其中si
是update_scatter_index
中的各个元素,如果index_vector_dim
<rank(scatter_indices)
,则会在index_vector_dim
索引处插入:
。- 否则为
[scatter_indices[update_scatter_index]]
。
- 对于
axes(inputs[0])
中的d_input
,- 如果
d_input = scatter_dims_to_operand_dims[d_start]
,则为full_start_index[d_input] = start_index[d_start]
。 - 否则为
full_start_index[d_input] = 0
。
- 如果
update_window_index = update_index[update_window_dims...]
.full_window_index = [wi0, ..., 0, ..., wiN]
,其中wi
是update_window_index
中的各个元素,0
是在inserted_window_dims
中的索引处插入。result_index = full_start_index + full_window_index
.
鉴于此,results = exec(schedule, inputs)
,其中:
schedule
是实现定义的index_space(updates[0])
排列。exec([update_index, ...], results) = exec([...], updated_results)
,其中:- 如果
result_index
在shape(results...)
的边界内 updates_converted = to_destination_type( updates...[update_index], type(func_inputs(update_computation) [len(func_inputs(update_computation))//2:])... )
updated_values = update_computation(results...[result_index], updates_converted)
updated_results
是results
的副本,其中results...[result_index]
设置为updated_values...
。- 否则
updated_results = results
.
- 如果
exec([], results) = results
.
如果 indices_are_sorted
为 true
,则实现可以假定 scatter_indices
相对于 scatter_dims_to_operand_dims
进行排序,否则行为将处于未定义状态。更正式地说,对于 indices(result)
中的所有 i1 < i2
,full_start_index(i1)
<= full_start_index(i2)
。
如果 unique_indices
为 true
,则实现可以假定分散的所有 result_index
索引都是唯一的。如果 unique_indices
为 true
,但分散到的索引不是唯一的,则行为未定义。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | inputs |
可变张量数量或每个张量量化张量 | (C1)、(C2)、(C4-C6)、(C10)、(C13)、(C15-C16) |
(I2) | scatter_indices |
整数类型的张量 | (C4)、(C11)、(C14) |
(I3) | updates |
可变张量数量或每个张量量化张量 | (C3-C6)、(C8) |
(I4) | update_window_dims |
si64 类型的一维张量常数 |
(C2)、(C4)、(C7)、(C8) |
(I5) | inserted_window_dims |
si64 类型的一维张量常数 |
(C2)、(C4)、(C9)、(C10) |
(I6) | scatter_dims_to_operand_dims |
si64 类型的一维张量常数 |
(C11-C13) |
(I7) | index_vector_dim |
si64 类型的常量 |
(C4)、(C11)、(C14) |
(I8) | indices_are_sorted |
i1 类型的常量 |
|
(I9) | unique_indices |
i1 类型的常量 |
|
(I10) | update_computation |
function | (C15) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变张量数量或每个张量量化张量 | (C15-C17) |
限制条件
- (C1)
same(shape(inputs...))
。 - (C2)
rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims)
。 - (C3)
same(shape(updates...))
。 - (C4)
shape(updates[0]) = combine(update_scatter_dim_sizes, update_window_dim_sizes)
,其中:update_scatter_dim_sizes = shape(scatter_indices)
,只不过不包含与index_vector_dim
对应的scatter_indices
的尺寸大小。update_window_dim_sizes <= shape(inputs[0])
,只不过inputs[0]
中与inserted_window_dims
对应的尺寸尺寸不包括在内。combine
将update_scatter_dim_sizes
放置在与update_scatter_dims
对应的轴上,将update_window_dim_sizes
放置在与update_window_dims
对应的轴上。
- (C5)
0 < size(inputs) = size(updates) = N
。 - (C6)
element_type(updates...) = element_type(inputs...)
。 - (C7)
is_unique(update_window_dims) and is_sorted(update_window_dims)
。 - (C8)
0 <= update_window_dims < rank(updates[0])
。 - (C9)
is_unique(inserted_window_dims) and is_sorted(update_window_dims)
。 - (C10)
0 <= inserted_window_dims < rank(inputs[0])
。 - (C11)
size(scatter_dims_to_operand_dims) = index_vector_dim < rank(scatter_indices) ? dim(scatter_indices, index_vector_dim) : 1
。 - (C12)
is_unique(scatter_dims_to_operand_dims)
。 - (C13)
0 <= scatter_dims_to_operand_dims < rank(inputs[0])
。 - (C14)
0 <= index_vector_dim <= rank(scatter_indices)
。 - (C15)
update_computation
的类型为(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)
,其中is_promotable(element_type(inputs[i]), Ei)
。 - (C16)
shape(inputs...) = shape(results...)
。 - (C17) 对
[0,N)
中的所有i
执行element_type(results[i]) = Ei
。
示例
// %input: [
// [[1, 2], [3, 4], [5, 6], [7, 8]],
// [[9, 10], [11, 12], [13, 14], [15, 16]],
// [[17, 18], [19, 20], [21, 22], [23, 24]]
// ]
// %scatter_indices: [[[0, 2], [1, 0], [2, 1]], [[0, 1], [1, 0], [0, 9]]]
// %update: [
// [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]],
// [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]]
// ]
%result = "stablehlo.scatter"(%input, %scatter_indices, %update) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
scatter_dimension_numbers = #stablehlo.scatter<
update_window_dims = [2, 3],
inserted_window_dims = [0],
scatter_dims_to_operand_dims = [1, 0],
index_vector_dim = 2>,
indices_are_sorted = false,
unique_indices = false
} : (tensor<3x4x2xi64>, tensor<2x3x2xi64>, tensor<2x3x2x2xi64>) -> tensor<3x4x2xi64>
// %result: [
// [[1, 2], [5, 6], [7, 8], [7, 8]],
// [[10, 11], [12, 13], [14, 15], [16, 17]],
// [[18, 19], [20, 21], [21, 22], [23, 24]]
// ]
选择
语义
生成一个 result
张量,其中根据 pred
中相应元素的值从 on_true
或 on_false
张量中选择每个元素。更正式地说,是 result[result_index] = pred_element ? on_true[result_index] :
on_false[result_index]
,其中 pred_element = rank(pred) = 0 ? pred[] :
pred[result_index]
。对于量化类型,请执行 dequantize_select_quantize(pred, on_true, on_false, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | pred |
i1 类型的张量 |
(C1) |
(I2) | on_true |
或每张量量化张量 | (C1-C2) |
(I3) | on_false |
或每张量量化张量 | (C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C2) |
限制条件
- (C1)
rank(pred) = 0 or shape(pred) = shape(on_true)
。 - (C2)
baseline_type(on_true) = baseline_type(on_false) = baseline_type(result)
。
示例
// %pred: [[false, true], [true, false]]
// %on_true: [[1, 2], [3, 4]]
// %on_false: [[5, 6], [7, 8]]
%result = "stablehlo.select"(%pred, %on_true, %on_false) : (tensor<2x2xi1>, tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 2], [3, 8]]
select_and_scatter
语义
根据使用 select
的 input
张量的 reduce_window
的结果,使用 scatter
散布来自 source
张量的值,并生成一个 result
张量。
下图显示了如何使用一个具体示例根据 operand
和 source
计算 result
中的元素。
更正式一点:
输入以下内容的
selected_values = reduce_window_without_init(...)
:- `inputs = [操作数]。
window_dimensions
、window_strides
和padding
按原样使用。base_dilations = windows_dilations = 1
.body
的定义如下:
def body(arg0: tensor<E>, arg1: tensor<E>) -> tensor<E>: return select(arg0, arg1) ? arg0 : arg1;
其中,
E = element_type(operand)
和reduce_window_without_init
的工作方式与reduce_window
完全相同,只不过底层reduce
的schedule
(请参阅 reduce)不包括 init 值。目前未指定如果相应窗口没有值时会发生什么情况 (#731)。result[result_index] = reduce([source_values], [init_value], [0], scatter)
,其中:source_values = [source[source_index] for source_index in source_indices]
.- 如果
selected_values[source_index]
具有operand_index
中的operand
元素,则为selected_index(source_index) = operand_index
。 source_indices = [source_index for source_index in indices(source) if selected_index(source_index) = result_index]
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1-C4)、(C6)、(C8-C11) |
(I2) | source |
或每张量量化张量 | (C1)、(C2) |
(I3) | init_value |
0 维张量或每张量量化张量 | (C3) |
(I4) | window_dimensions |
si64 类型的一维张量常数 |
(C2)、(C4)、(C5) |
(I5) | window_strides |
si64 类型的一维张量常数 |
(C2)、(C6)、(C7) |
(I6) | padding |
si64 类型的二维张量常数 |
(C2)、(C8) |
(I7) | select |
function | (C9) |
(I8) | scatter |
function | (C10) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C11-C12) |
限制条件
- (C1)
element_type(operand) = element_type(source)
。 - (C2)
shape(source) = num_windows
,其中:padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1]
.is_empty_window = padded_operand_shape = 0 || window_dimensions > padded_operand_shape
.num_windows = is_empty_window ? 0 : floor((padded_operand_shape - window_dimensions) / window_strides) + 1
.
- (C3)
element_type(init_value) = element_type(operand)
。 - (C4)
size(window_dimensions) = rank(operand)
。 - (C5)
0 < window_dimensions
。 - (C6)
size(window_strides) = rank(operand)
。 - (C7)
0 < window_strides
。 - (C8)
shape(padding) = [rank(operand), 2]
。 - (C9)
select
的类型为(tensor<E>, tensor<E>) -> tensor<i1>
,其中E = element_type(operand)
。 - (C10)
scatter
的类型为(tensor<E>, tensor<E>) -> tensor<E>
,其中is_promotable(element_type(operand), E)
。 - (C11)
shape(operand) = shape(result)
。 - (C12)
element_type(result) = E
。
示例
// %operand: [[1, 5], [2, 5], [3, 6], [4, 4]]
// %source: [[5, 6], [7, 8]]
// %init_value: 0
%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.compare"(%arg0, %arg1) {
comparison_direction = #stablehlo<comparison_direction GE>
} : (tensor<i64>, tensor<i64>) -> tensor<i1>
"stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
window_dimensions = dense<[3, 1]> : tensor<2xi64>,
window_strides = dense<[2, 1]> : tensor<2xi64>,
padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi64>, tensor<2x2xi64>, tensor<i64>) -> tensor<4x2xi64>
// %result: [[0, 0], [0, 0], [5, 14], [7, 0]]
send
语义
将 inputs
发送到通道 channel_id
并生成一个 result
令牌。
如果 is_host_transfer
为 true
,则操作会将数据转移到主机。否则,它会将数据传输到其他设备。这意味着,是由实现定义的。该标志与 channel_type
中提供的信息重复,因此将来我们计划只保留其中一个 (#666)。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | inputs |
可变数量的张量或量化张量 | |
(I2) | token |
token |
|
(I3) | channel_id |
si64 类型的常量 |
|
(I4) | channel_type |
DEVICE_TO_DEVICE 和 DEVICE_TO_HOST 的枚举 |
(C1) |
(I5) | is_host_transfer |
i1 类型的常量 |
(C1) |
输出
名称 | 类型 |
---|---|
result |
token |
限制条件
- (C1)
channel_type
的定义如下:- 如果
is_host_transfer = true
,则为DEVICE_TO_HOST
; - 否则为
DEVICE_TO_DEVICE
。
- 如果
示例
%result = "stablehlo.send"(%operand, %token) {
channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>,
is_host_transfer = true
} : (tensor<2x2xi64>, !stablehlo.token) -> !stablehlo.token
shift_left
语义
对 lhs
张量执行 rhs
位数的元素级左移运算,并生成一个 result
张量。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
整数类型的张量 | (C1) |
(I2) | rhs |
整数类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数类型的张量 | (C1) |
限制条件
- (C1)
type(lhs) = type(rhs) = type(result)
。
示例
// %lhs: [-1, 0, 1]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_left"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-2, 0, 8]
shift_right_arithmetic
语义
对 lhs
张量执行 rhs
位数的元素级算术右移运算,并生成一个 result
张量。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
整数类型的张量 | (C1) |
(I2) | rhs |
整数类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数类型的张量 | (C1) |
限制条件
- (C1)
type(lhs) = type(rhs) = type(result)
。
示例
// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_arithmetic"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-1, 0, 1]
shift_right_logical
语义
对 lhs
张量执行 rhs
位数的元素级逻辑右移运算,并生成一个 result
张量。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
整数类型的张量 | (C1) |
(I2) | rhs |
整数类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数类型的张量 | (C1) |
限制条件
- (C1)
type(lhs) = type(rhs) = type(result)
。
示例
// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_logical"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [9223372036854775807, 0, 1]
签名
语义
按元素返回 operand
的符号,并生成一个 result
张量。更正式地说,对于每个 x
元素,可以使用 Python 语法表示语义,如下所示:
def sign(x):
if is_integer(x):
if compare(x, 0, LT, SIGNED): return -1
if compare(x, 0, EQ, SIGNED): return 0
return 1
elif is_float(x):
if is_nan(x): return NaN
if compare(x, -0.0, EQ, FLOAT): return -0.0
if compare(x, +0.0, EQ, FLOAT): return +0.0
if compare(x, 0.0, LT, FLOAT): return -1.0
return 1.0
elif is_complex(x):
if is_nan(real(x)) or is_nan(imag(x)): return (NaN, NaN)
if compare(x, (0.0, 0.0), EQ, FLOAT): return (0.0, 0.0)
return divide(x, convert(abs(x), type(x)))
对于量化类型,请执行 dequantize_op_quantize(sign, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
带符号整数、浮点或复杂类型的张量,或每个张量量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
带符号整数、浮点或复杂类型的张量,或每个张量量化张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// operand: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]
%result = "stablehlo.sign"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// %result: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]
正弦
语义
对 operand
张量执行元素级正弦运算并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
sin
。 - 对于复数:复数正弦。
- 对于量化类型:
dequantize_op_quantize(sine, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [
// [0.0, 1.57079632], // [0, pi/2]
// [3.14159265, 4.71238898] // [pi, 3pi/2]
// ]
%result = "stablehlo.sine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [0.0, -1.0]]
slice
语义
使用静态计算的起始索引从 operand
中提取 Slice,并生成 result
张量。start_indices
包含每个维度切片的起始索引,limit_indices
包含每个维度切片的结束索引(不含该索引),strides
包含每个维度的步长。
更正式地说,是 result[result_index] = operand[operand_index]
,其中 operand_index = start_indices + result_index * strides
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或每张量量化张量 | (C1-C3)、(C5) |
(I2) | start_indices |
si64 类型的一维张量常数 |
(C2)、(C3)、(C5) |
(I3) | limit_indices |
si64 类型的一维张量常数 |
(C2)、(C3)、(C5) |
(I4) | strides |
si64 类型的一维张量常数 |
(C2)、(C4) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或每张量量化张量 | (C1)、(C5) |
限制条件
- (C1)
element_type(operand) = element_type(result)
。 - (C2)
size(start_indices) = size(limit_indices) = size(strides) = rank(operand)
。 - (C3)
0 <= start_indices <= limit_indices <= shape(operand)
。 - (C4)
0 < strides
。 - (C5)
shape(result) = ceil((limit_indices - start_indices) / strides)
。
示例
// %operand: [
// [0, 0, 0, 0],
// [0, 0, 1, 1],
// [0, 0, 1, 1]
// ]
%result = "stablehlo.slice"(%operand) {
start_indices = dense<[1, 2]> : tensor<2xi64>,
limit_indices = dense<[3, 4]> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>
// % result: [
// [1, 1],
// [1, 1]
// ]
sort
语义
根据 comparator
,将沿 dimension
维度的 inputs
的一维切片一起排序,并生成 results
。
与其他运算中的类似输入不同,dimension
允许负值,具有下述语义。将来,出于一致性原因,我们可能会禁止使用此参数 (#1377)。
如果 is_stable
为 true,则排序是稳定的,即保留比较器认为相等的元素的相对顺序。对于只有一个输入的情况,当且仅当 comparator(e1, e2) = comparator(e2, e1) = false
时,比较器才会认为两个元素 e1
和 e2
相等。如需了解其如何泛化到多个输入,请参阅以下格式。
更正式地说,对于 index_space(results[0])
中的所有 result_index
:
adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension
.result_slice = [ri0, ..., :, ..., riR-1]
,其中riN
是result_index
中的各个元素,:
在adjusted_dimension
处插入。inputs_together = (inputs[0]..., ..., inputs[N-1]...)
.results_together[result_slice] = sort(inputs_together[result_slice], comparator_together)
.- 其中
sort
按非降序对一维切片进行排序,预期如果左侧参数小于右侧第二个参数,comparator_together
会返回true
。 def comparator_together(lhs_together, rhs_together): args = [] for (lhs_el, rhs_el) in zip(lhs_together, rhs_together): args.append(lhs_el) args.append(rhs_el) return comparator(*args)
(results[0]..., ..., results[N-1]...) = results_together
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | inputs |
可变张量数量或每个张量量化张量 | (C1-C5) |
(I2) | dimension |
si64 类型的常量 |
(C4) |
(I3) | is_stable |
i1 类型的常量 |
|
(I4) | comparator |
function | (C5) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变张量数量或每个张量量化张量 | (C2)、(C3) |
限制条件
- (C1)
0 < size(inputs)
。 - (C2)
type(inputs...) = type(results...)
。 - (C3)
same(shape(inputs...) + shape(results...))
。 - (C4)
-R <= dimension < R
,其中R = rank(inputs[0])
。 - (C5)
comparator
的类型为(tensor<E1>, tensor<E1>, ..., tensor<EN-1>, tensor<EN-1>) -> tensor<i1>
,其中Ei = element_type(inputs[i])
。
示例
// %input0 = [[1, 2, 3], [3, 2, 1]]
// %input1 = [[3, 2, 1], [1, 2, 3]]
%result0, %result1 = "stablehlo.sort"(%input0, %input1) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>, %arg2: tensor<i64>, %arg3: tensor<i64>):
%predicate = "stablehlo.compare"(%arg0, %arg1) {
comparison_direction = #stablehlo<comparison_direction GT>
} : (tensor<i64>, tensor<i64>) -> tensor<i1>
"stablehlo.return"(%predicate) : (tensor<i1>) -> ()
}) {
dimension = 0 : i64,
is_stable = true
} : (tensor<2x3xi64>, tensor<2x3xi64>) -> (tensor<2x3xi64>, tensor<2x3xi64>)
// %result0 = [[3, 2, 3], [1, 2, 1]]
// %result1 = [[1, 2, 1], [3, 2, 3]]
sqrt
语义
对 operand
张量执行元素级平方根运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
squareRoot
。 - 对于复数:复数平方根。
- 对于量化类型:
dequantize_op_quantize(sqrt, operand, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [[0.0, 1.0], [4.0, 9.0]]
%result = "stablehlo.sqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [2.0, 3.0]]
subtract
语义
对 lhs
和 rhs
两个张量执行元素级减法,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于整数:整数减法。
- 对于浮点数:IEEE-754 中的
subtraction
。 - 对于复数:复数减法。
- 对于量化类型:
dequantize_op_quantize(subtract, lhs, rhs, type(result))
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
(I2) | rhs |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
整数、浮点或复杂类型的张量,或每个张量的量化张量 | (C1) |
限制条件
- (C1)
baseline_type(lhs) = baseline_type(rhs) = baseline_type(result)
。
示例
// %lhs: [[6, 8], [10, 12]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.subtract"(%lhs, %rhs) : (tensor<2x2xf32>, tensor<2x2xf32>) -> (tensor<2x2xf32>)
// %result: [[1, 2], [3, 4]]
双色
语义
对 operand
张量执行元素级双曲正切运算,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于浮点数:IEEE-754 中的
tanh
。 - 对于复数:复双曲正切。
- 对于量化类型:
dequantize_op_quantize(tanh, operand, type(result))
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_type(operand) = baseline_type(result)
。
示例
// %operand: [-1.0, 0.0, 1.0]
%result = "stablehlo.tanh"(%operand) : (tensor<3xf32>) -> tensor<3xf32>
// %result: [-0.76159416, 0.0, 0.76159416]
转置
语义
使用 permutation
排列 operand
张量的维,并生成一个 result
张量。更正式地说,是 result[result_index] = operand[operand_index]
,其中 result_index[d] = operand_index[permutation[d]]
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
或量化张量 | (C1-C4) |
(I2) | permutation |
si64 类型的一维张量常数 |
(C2-C4) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
或量化张量 | (C1)、(C3-C4) |
限制条件
- (C1)
element_type(result)
的计算公式如下:element_type(operand)
(如果!is_per_axis_quantized(operand)
)。element_type(operand)
,但quantization_dimension(operand)
和quantization_dimension(result)
可能会有所不同。
- (C2)
permutation
是range(rank(operand))
的排列。 - (C3)
shape(result) = dim(operand, permutation...)
。 - (C4) 如果为
is_per_axis_quantized(result)
,则设为quantization_dimension(operand) = permutation(quantization_dimension(result))
。
示例
// %operand: [
// [[1,2], [3,4], [5,6]],
// [[7,8], [9,10], [11,12]]
// ]
%result = "stablehlo.transpose"(%operand) {
permutation = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
// %result: [
// [[1,7], [3,9], [5,11]],
// [[2,8], [4,10], [6,12]]
// ]
triangular_solve
语义
求解具有下层或上三角系数矩阵的线性方程组的批次。
更正式地说,给定 a
和 b
,当 left_side
为 true
或 x * op(a[i0, ..., iR-3, :, :]) = b[i0, ..., iR-3, :, :]
为 false
时,result[i0, ..., iR-3, :, :]
是 op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :]
的解,求出变量 x
,其中 op(a)
由 transpose_a
决定,可以是以下几项之一:left_side
NO_TRANSPOSE
:按原样使用a
执行操作。TRANSPOSE
:对a
的转置执行操作。ADJOINT
:对a
的共置转置执行操作。
如果 lower
为 true
或 a
的上三角形,则仅从 a
的下三角形读取输入数据。输出数据在同一三角形中返回;另一个三角形中的值由实现定义。
如果 unit_diagonal
为 true,则实现可以假定 a
的对角线元素等于 1,否则行为将处于未定义状态。
对于量化类型,请执行 dequantize_op_quantize(lambda x, y: triangular_solve(x, y, left_side, lower,
unit_diagonal, transpose_a), a, b, type(result))
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | a |
浮点或复杂类型或每个张量量化张量的张量 | (C1-C3) |
(I2) | b |
浮点或复杂类型或每个张量量化张量的张量 | (C1-C4) |
(I3) | left_side |
i1 类型的常量 |
(C3) |
(I4) | lower |
i1 类型的常量 |
|
(I5) | unit_diagonal |
i1 类型的常量 |
|
(I6) | transpose_a |
NO_TRANSPOSE 、TRANSPOSE 和 ADJOINT 的枚举 |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点或复杂类型或每个张量量化张量的张量 | (C1) |
限制条件
- (C1)
baseline_element_type(a) = baseline_element_type(b)
。 - (C2)
2 <= rank(a) = rank(b) = R
。 - (C3)
shape(a)
和shape(b)
之间的关系定义如下:shape(a)[:-3] = shape(b)[:-3]
.dim(a, -2) = dim(a, -1) = dim(b, left_side ? -2 : -1)
.
- (C4)
baseline_type(b) = baseline_type(result)
。
示例
// %a = [
// [1.0, 0.0, 0.0],
// [2.0, 4.0, 0.0],
// [3.0, 5.0, 6.0]
// ]
// %b = [
// [2.0, 0.0, 0.0],
// [4.0, 8.0, 0.0],
// [6.0, 10.0, 12.0]
// ]
%result = "stablehlo.triangular_solve"(%a, %b) {
left_side = true,
lower = true,
unit_diagonal = false,
transpose_a = #stablehlo<transpose NO_TRANSPOSE>
} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// %result: [
// [2.0, 0.0, 0.0],
// [0.0, 2.0, 0.0],
// [0.0, 0.0, 2.0]
// ]
tuple
语义
根据值 val
生成 result
元组。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | val |
值数量不等 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
tuple | (C1) |
限制条件
- (C1)
result
的类型为tuple<E0, ..., EN-1>
,其中Ei = type(val[i])
。
示例
// %val0: [1.0, 2.0]
// %val1: (3)
%result = "stablehlo.tuple"(%val0, %val1) : (tensor<2xf32>, tuple<tensor<i32>>) -> tuple<tensor<2xf32>, tuple<tensor<i32>>>
// %result: ([1.0, 2.0], (3))
uniform_dequantize
语义
根据 operand
类型定义的量化参数,对量化张量 operand
执行元素级转换,使其成为浮点张量 result
。
更正式地说,result = dequantize(operand)
。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
量化张量 | (C1)、(C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
浮点类型的张量 | (C1)、(C2) |
限制条件
- (C1)
shape(operand) = shape(result)
。 - (C2)
element_type(result) = expressed_type(operand)
。
示例
// %operand: [10, 10]
%result = "stablehlo.uniform_dequantize"(%operand) : (tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>) -> tensor<2xf32>
// %result: [4.0, 15.0]
uniform_quantize
语义
根据 result
类型定义的量化参数,将浮点张量或量化张量 operand
转换为量化张量 result
。
更正式地说,
- 如果为
is_float(operand)
:result = quantize(operand, type(result))
.
- 如果为
is_quantized(operand)
:float_result = dequantize(operand)
.result = quantize(float_result, type(result))
.
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
浮点或量化类型的张量 | (C1)、(C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
量化张量 | (C1)、(C2) |
限制条件
- (C1)
shape(operand) = shape(result)
。 - (C2)
expressed_type(result) = is_float(operand) ? element_type(operand) : expressed_type(operand)
。
示例
// %operand: [4.0, 15.0]
%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>
// %result: [10, 10]
// %operand: [10, 10]
%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>) -> tensor<2x!quant.uniform<i8:f32:0, {0.1:-20,0.2:-30}>>
// %result: [20, 45]
而
语义
执行 body
函数 0 次或多次时产生输出,同时 cond
函数输出 true
。更正式地说,可以使用 Python 语法表示语义,如下所示:
internal_state = operand
while cond(*internal_state):
internal_state = body(*internal_state)
results = internal_state
无限循环的行为待定 (#383)。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | operand |
可变数量的张量、量化张量或词元 | (C1-C3) |
(I2) | cond |
function | (C1) |
(I3) | body |
function | (C2) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
results |
可变数量的张量、量化张量或词元 | (C3) |
限制条件
- (C1)
cond
的类型为(T0, ..., TN-1) -> tensor<i1>
,其中Ti = type(operand[i])
。 - (C2)
body
的类型为(T0, ..., TN-1) -> (T0, ..., TN-1)
,其中Ti = type(operand[i])
。 - (C3)
type(results...) = type(operand...)
。
示例
// %init_i: 1
// %init_sum: 0
// %one: 1
// %ten: 10
%results0, %results1 = "stablehlo.while"(%init_i, %init_sum) ({
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%cond = "stablehlo.compare"(%arg0, %ten) {
comparison_direction = #stablehlo<comparison_direction LT>
} : (tensor<i64>, tensor<i64>) -> tensor<i1>
stablehlo.return %cond : tensor<i1>
}, {
^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
%new_sum = stablehlo.add %arg1, %one : tensor<i64>
%new_i = stablehlo.add %arg0, %one : tensor<i64>
stablehlo.return %new_i, %new_sum : tensor<i64>, tensor<i64>
}) : (tensor<i64>, tensor<i64>) -> (tensor<i64>, tensor<i64>)
// %results0: 10
// %results1: 10
异或
语义
对两个张量 lhs
和 rhs
执行元素级 XOR,并生成一个 result
张量。根据元素类型,执行以下操作:
- 对于布尔值:逻辑 XOR。
- 对于整数:按位 XOR。
输入内容
标签 | 名称 | 类型 | 限制条件 |
---|---|---|---|
(I1) | lhs |
布尔值或整数类型的张量 | (C1) |
(I2) | rhs |
布尔值或整数类型的张量 | (C1) |
输出
名称 | 类型 | 限制条件 |
---|---|---|
result |
布尔值或整数类型的张量 | (C1) |
限制条件
- (C1)
type(lhs) = type(rhs) = type(result)
。
示例
// Bitwise operation with with integer tensors
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.xor"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[4, 4], [4, 12]]
// Logical operation with with boolean tensors
// %lhs: [[false, false], [true, true]]
// %rhs: [[false, true], [false, true]]
%result = "stablehlo.xor"(%lhs, %rhs) : (tensor<2x2xi1>, tensor<2x2xi1>) -> tensor<2x2xi1>
// %result: [[false, true], [true, false]]
执行
顺序执行
StableHLO 程序通过向 main
函数提供输入值并计算输出值来执行。函数的输出值通过执行根于相应 return
操作的操作图来计算。
只要执行顺序与数据流保持一致(即操作是否在使用之前执行),则执行顺序就是由实现定义的。在 StableHLO 中,所有附带效应的操作都会消耗一个令牌,并生成一个令牌(多个令牌可通过 after_all
多路复用为一个令牌),因此附带效应的执行顺序也与数据流一致。上述示例程序的可能执行顺序为 %0
→ %1
→ %2
→ %3
→ %4
→ return
或 %3
→ %0
→ %1
→ %2
→ %4
→ return
。
更正式地说,StableHLO 进程是以下项的组合:1) StableHLO 程序,2) 操作状态(尚未执行,已经执行),3) 进程正在处理的中间值。该过程从 main
函数的输入值开始,一直经历更新操作状态和中间值的操作图,最后以输出值结束。有关进一步正式的信息,请参阅待定 (#484)。
并行执行
StableHLO 程序可以并行执行,并由 num_partitions
组织到 num_replicas
的 2D 进程网格中,二者的类型均为 ui32
。
在 StableHLO 进程网格中,系统会同时执行 StableHLO 进程的 num_replicas * num_partitions
。每个进程都有一个唯一的 process_id = (replica_id, partition_id)
,其中 replica_ids = range(num_replicas)
中的 replica_id
和 partition_ids = range(num_partitions)
中的 partition_id
都具有 ui32
类型。
对于每个程序,进程网格的大小都是静态已知的(未来,我们计划使其成为 StableHLO 程序的显式部分 #650),而进程网格中的位置对于每个进程都是静态已知的。每个进程都可以通过 replica_id
和 partition_id
操作访问它在进程网格中的位置。
在进程网格中,程序可以全部相同(采用“Single Program, multiple Data”样式)、不同(在“Multiple Program,Multiple Data”样式)中,也可以介于两者之间。未来,我们计划支持定义并行 StableHLO 程序的其他习语,包括 GSPMD (#619)。
在进程网格中,进程大多彼此独立 - 它们具有单独的操作状态,单独的输入/中间/输出值,并且大多数操作在进程之间单独执行,但下面介绍的少数集合操作除外。
鉴于大多数操作的执行仅使用来自同一进程的值,因此按名称引用这些值通常没有歧义。但是,在描述集合操作的语义时,这样做还不够,这就导致了 name@process_id
表示法在特定进程中引用值 name
。(从这个角度来看,非限定的 name
可以看作是 name@(replica_id(), partition_id())
的简写形式。)
所有进程的执行顺序都是由实现定义的,但由点对点通信和集体操作引入的同步除外(如下所述)。
点对点通信
StableHLO 进程可以通过 StableHLO 通道相互通信。渠道由 si64
类型的正 ID 表示。通过各种操作,您可以向通道发送值,然后从通道接收值。
进一步规范化,例如,这些频道 ID 的来源、进程如何发现它们以及它们会引入哪种同步,目前尚处于待定状态 (#484)。
流式通信
每个 StableHLO 进程都可以访问两个流接口:
- 可供读取的信息流广告。
- 可写入的外馈。
与用于在进程之间进行通信(因此两端都有进程)的渠道不同,infeed 和 outfeed 的另一端实现则由其定义。
要进一步规范化,例如流式通信如何影响执行顺序以及会引入何种同步,目前尚待确定 (#484)。
集体行动
StableHLO 有六个共同操作:all_gather
、all_reduce
、all_to_all
、collective_broadcast
、collective_permute
和 reduce_scatter
。所有这些操作都会将 StableHLO 进程网格中的进程拆分为 StableHLO 进程组,并在每个进程组(独立于其他进程组)中执行联合计算。
在每个进程组内,共同操作可能会引入同步屏障。待进一步规范化,例如,详细阐述此同步的确切发生时间、进程究竟如何到达此屏障,以及未发生时会发生什么情况,则属于待定 (#484)。
如果进程组涉及跨分区通信,即进程组中一些进程的分区 ID 不同,则集合操作的执行需要一个通道,而集合操作的执行必须提供一个类型为 si64
的正 channel_id
。跨副本通信不需要通道。
这些共同操作执行的计算特定于单个操作,并在上面的各个操作部分中进行了介绍。不过,将进程网格拆分为进程组的策略会在这些操作之间共享,本部分将对此进行介绍。从更正式的角度来说,StableHLO 支持以下四种策略。
cross_replica
每个进程组中仅发生跨副本通信。此策略采用 replica_groups
(副本 ID 列表)作为参数,并按 partition_ids
计算 replica_groups
的笛卡尔积。replica_groups
必须具有唯一的元素,并覆盖所有 replica_ids
。更正式地说,使用 Python 语法:
def cross_replica(replica_groups: List[List[ReplicaId]]) -> List[List[ProcessId]]:
for replica_group in replica_groups:
for partition_id in partition_ids:
process_group = []
for replica_id in replica_group:
process_group.append((replica_id, partition_id))
yield process_group
例如,对于 replica_groups = [[0, 1], [2, 3]]
和 num_partitions = 2
,cross_replica
将生成 [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(2, 0), (3, 0)], [(2, 1), (3, 1)]]
。
cross_partition
每个进程组中仅发生跨分区通信。此策略采用 partition_groups
(分区 ID 列表)作为参数,并按 replica_ids
计算 partition_groups
的笛卡尔积。partition_groups
必须具有唯一的元素,并覆盖所有 partition_ids
。更正式地说,使用 Python 语法:
def cross_partition(partition_groups: List[List[PartitionId]]) -> List[List[ProcessId]]:
for partition_group in partition_groups:
for replica_id in replica_ids:
process_group = []
for partition_id in partition_group:
process_group.append((replica_id, partition_id))
yield process_group
例如,对于 partition_groups = [[0, 1]]
和 num_replicas = 4
,cross_partition
将生成 [[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)], [(3, 0), (3, 1)]]
。
cross_replica_and_partition
跨副本通信和跨分区通信都可能在每个进程组内进行。此策略采用 replica_groups
(副本 ID 列表)作为参数,并按 partition_ids
计算每个 replica_group
的笛卡尔积。replica_groups
必须具有唯一的元素,并覆盖所有 replica_ids
。更正式地说,使用 Python 语法:
def cross_replica_and_partition(replica_groups: List[List[ReplicaId]]) -> List[List[ProcessId]]:
for replica_group in replica_groups:
process_group = []
for partition_id in partition_ids:
for replica_id in replica_group:
process_group.append((replica_id, partition_id))
yield process_group
例如,对于 replica_groups = [[0, 1], [2, 3]]
和 num_partitions = 2
,cross_replica_and_partition
将生成 [[(0, 0), (1, 0), (0, 1), (1, 1)], [(2, 0), (3, 0), (2, 1), (3, 1)]]
。
flattened_ids
此策略采用 flattened_id_groups
(replica_id * num_partitions + partition_id
形式的“扁平化”进程 ID 列表)并将其转换为进程 ID。flattened_id_groups
必须具有唯一的元素,并覆盖所有 process_ids
。更正式地说,使用 Python 语法:
def flattened_ids(flattened_id_groups: List[List[ui32]]) -> List[List[ProcessId]]:
for flattened_id_group in flattened_id_groups:
process_group = []
for flattened_id in flattened_id_group:
replica_id = flattened_id // num_partitions
partition_id = flattened_id % num_partitions
process_group.append((replica_id, partition_id))
yield process_group
例如,对于 flattened_id_groups = [[0, 1, 2, 3], [4, 5, 6, 7]]
、num_replicas = 4
和 num_partitions = 2
,flattened_ids
将生成 [[(0, 0), (0, 1), (1, 0), (1, 1)], [(2, 0), (2, 1), (3, 0), (3, 1)]]
。
准确率
目前,StableHLO 无法保证数值准确性,但将来可能会发生变化 (#1156)。
错误数
StableHLO 程序通过针对单个操作的一系列广泛限制条件进行验证,从而在运行时间之前排除许多类别的错误。不过,仍然可能会出现错误情况,例如发生整数溢出、出界访问等。除非明确调用,否则所有这些错误都会导致实现定义的行为,但将来可能会发生变化 (#1157)。
此规则的一个例外情况是,StableHLO 程序中的浮点异常具有明确定义的行为。导致 IEEE-754 标准定义的异常(无效操作、除零、上溢、下溢或不精确异常)的操作会生成默认结果(如标准中所定义),并在不引发相应状态标记的情况下继续执行;类似于标准中的 raiseNoFlag
异常处理。非标准运算(例如复杂算术和某些先验函数)的例外情况是由实现定义的。
Notation
为了描述语法,本文档使用了经过修改的 EBNF 语法的 ISO 变种(ISO/IEC 14977:1996、Wikipedia),并进行了两项修改:1) 使用 ::=
而不是 =
定义规则;
2) 串联使用并置(而不是 ,
)来表示。
为了描述语义(即在“类型”“常量”和“操作”部分中),我们使用的公式基于扩展的 Python 语法,并支持简明表达数组运算(如下所述)。这非常适合小的代码段,但在极少数情况下,当需要较大的代码段时,我们会使用始终明确引入的 vanilla Python 语法。
公式
我们基于 dot_general
规范中的示例来探索公式的工作原理。此操作的一项限制如下所示:dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...)
。
此公式中使用的名称有两个来源:1) 全局函数,即 dim
、2) 相应程序元素的成员定义,即 dot_general
的“输入源”部分中定义的 lhs
、lhs_batching_dimensions
、rhs
和 rhs_batching_dimensions
输入。
如上所述,此公式的语法基于 Python,并且包含一些面向简洁的扩展程序。为了理解这个公式 让我们将其转换为原始的 Python 语法
A) 在这些公式中,我们使用 =
表示等式,因此获取 Python 语法的第一步是将 =
替换为 ==
,如下所示:dim(lhs, lhs_batching_dimensions...) == dim(rhs, rhs_batching_dimensions...)
。
B) 此外,这些公式支持省略号 (...
),后者可将标量表达式转换为张量表达式。简而言之,f(xs...)
大致表示“对于张量 xs
中的每个标量 x
,计算一个标量 f(x)
,然后将所有这些标量结果一起作为张量结果返回”。在原始 Python 语法中,我们的示例公式会变为:[dim(lhs, dim1) for dim1 in lhs_batching_dimensions] ==
[dim(rhs, dim2) for dim2 in rhs_batching_dimensions]
。
得益于省略号,通常可以避免在单个标量级别工作。不过,在某些棘手的情况下,可以像 gather
规范中的 start_indices[bi0, ..., :, ..., biN]
公式一样使用较低级别的半非正式语法。为了简明扼要,对于将此类语法转换为普通 Python 语言,我们不提供确切的形式,希望仍按照具体情况直观地理解该语法。
如果某些特定公式看起来不透明,请告知我们,我们将尽力改进。
此外,您会注意到,公式使用省略号扩展各种列表,包括张量、张量列表(例如,可能来自不同数量的张量)等。在这方面,我们无法提供确切的正式形式(例如,列表甚至都不是 StableHLO 类型系统的一部分),而是依赖于直观性系统来理解。
C) 我们使用的最后一种值得注意的记数方式是隐式广播。虽然 StableHLO 运算集不支持隐式广播,但公式还支持简洁性。简而言之,如果在需要张量的上下文中使用标量,该标量会被广播为预期的形状。
继续以 dot_general
为例,以下是另一个限制条件:0 <= lhs_batching_dimensions < rank(lhs)
。如 dot_general
规范中所定义,lhs_batching_dimensions
是一个张量,但是 0
和 rank(lhs)
都是标量。应用隐式广播后,公式将变为 [0, ..., 0] <= lhs_batching_dimensions < [rank(lhs), ..., rank(lhs)]
。
应用于特定 dot_general
运算时,此公式的计算结果为布尔值张量。将公式用作约束条件时,如果公式的计算结果为 true
或仅具有 true
元素的张量,则约束条件。
姓名
在公式中,词法范围包括:1) 全局函数,2) 成员定义,
3) 当地定义。下面提供了全局函数列表。元素定义列表取决于要应用该表示法的程序元素:
- 对于操作,成员定义包括“输入”和“输出”部分中引入的名称。
- 对于其他所有内容,成员定义包括程序元素的结构部分,以相应的 EBNF 非终端命名。大多数情况下,这些结构部分的名称是通过将非终端的名称转换为蛇形命名法(例如
IntegerLiteral
=>integer_literal
)来获取的,但有时名称会在进程中采用缩写(例如QuantizationStorageType
=>storage_type
),在这种情况下,明确引入名称,与“输入”/“输出”部分中的“输入”类似。 - 此外,成员定义始终包含
self
来引用相应的程序元素。
值
计算公式时,公式会使用以下类型的值:
1) Value
(实际值,例如 dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
;它们始终知道自己的类型);
2) Placeholder
(未来的值,例如 lhs
、rhs
或 result
;其实际值目前未知,只有其类型已知)
3) Type
(“类型”部分中定义的类型);
4) Function
(“全局函数”部分中定义的全局函数,如 lhs
、rhs
或 result
)
名称可能引用不同的值,具体取决于上下文。更具体地说,运算的“语义”部分(以及其他程序元素的等效部分)定义了运行时逻辑,因此所有输入都以 Value
的形式提供。相比之下,操作(以及等效项)的“约束”部分则定义了“编译时”逻辑,即通常在运行时之前执行的内容,因此只有常量输入可用作 Value
,其他输入仅可用作 Placeholder
。
姓名 | 在“Semantics”中 | 在“约束条件”中 |
---|---|---|
全局函数 | Function |
Function |
常量输入 | Value |
Value |
非常量输入 | Value |
Placeholder |
输出 | Value |
Placeholder |
本地定义 | 取决于定义 | 取决于定义 |
我们来看一个示例 transpose
操作:
%result = "stablehlo.transpose"(%operand) {
permutation = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
对于此操作,permutation
是一个常量,因此在语义和约束条件中,它都以 Value
的形式提供。相比之下,operand
和 result
可在语义中以 Value
的形式使用,但在约束条件中仅以 Placeholder
的形式提供。
函数
类型的构造
没有可用于构造类型的函数。相反,我们直接使用类型语法,因为它通常更简洁。例如,使用 (tensor<E>, tensor<E>) -> (tensor<E>)
,而不是 function_type(
[tensor_type([], E), tensor_type([], E)], [tensor_type([], E)])
。
类型函数
element_type
基于张量类型和量化张量类型定义,并分别返回相应TensorType
或QuantizedTensorType
的TensorElementType
或QuantizedTensorElementType
部分。
def element_type(x: Value | Placeholder | Type):
if type(x) == TensorType:
return tensor_element_type(x)
if type(x) == QuantizedTensorType:
return quantized_tensor_element_type(x)
if type(x) is not Type:
return element_type(type(x))
is_per_axis_quantized(x: Value | Placeholder | Type) -> Value
是is_quantized(x) and quantization_dimension(x) is not None
的快捷方式。is_per_tensor_quantized(x: Value | Placeholder | Type) -> Value
是is_quantized(x) and quantization_dimension(x) is None
的快捷方式。is_promotable(x: Type, y: Type) -> bool
检查类型x
是否可以提升为y
类型。当x
和y
为QuantizedTensorElementType
时,促销活动仅应用于storage_type
。这一特定版本的促销目前用于减少计算的上下文(如需了解详情,请参阅 RFC)。
def is_promotable(x: Type, y: Type) -> Value:
is_same_type = (is_bool(x) and is_bool(y)) or
(is_integer(x) and is_integer(y)) or (is_float(x) and is_float(y)) or
(is_complex(x) and is_complex(y)) or
(is_quantized(x) and is_quantized(y) and expressed_type(x) = expressed_type(y))
if is_same_type == False:
return False
if is_integer(x) or is_float(x):
return bitwidth(x) <= bitwidth(y)
if is_complex(x):
return bitwidth(element_type(x)) <= bitwidth(element_type(y))
if is_quantized(x):
return bitwidth(storage_type(x)) <= bitwidth(storage_type(y))
return false
is_quantized(x: Value | Placeholder | Type) -> Value
是is_quantized_tensor_element_type(x)
的快捷方式。is_type_name(x: Value | Placeholder | Type) -> Value
。适用于所有类型。例如,如果x
为FloatType
,则is_float(x)
会返回true
。如果x
是值或占位符,则此函数是is_type_name(type(x))
的快捷方式。max_value(x: Type) -> Value
返回TensorElementType
的最大值。如果x
不是TensorElementType
,则返回None
。min_value(x: Type) -> Value
会返回TensorElementType
的可能最小值。如果x
不是TensorElementType
,则返回None
。member_name(x: Value | Placeholder | Type) -> Any
。适用于所有类型的所有成员定义member_name
。例如,tensor_element_type(x)
会返回对应TensorType
的TensorElementType
部分。如果x
是值或占位符,则此函数是member_name(type(x))
的快捷方式。如果x
不是具有适当成员的类型,也不不是此类类型的值或占位符,则返回None
。
值的构造
operation_name(*xs: Value | Type) -> Value
。适用于所有操作。例如,add(lhs, rhs)
接受两个张量值lhs
和rhs
,并返回通过这些输入评估add
运算的输出。对于某些操作(例如broadcast_in_dim
),其输出类型为“承载”,即评估操作所需的输出。在本例中,该函数将这些类型作为参数。
值函数
Python 的所有运算符和函数都可用。例如,Python 中的订阅和切片表示法都可用于将索引编入张量、量化张量和元组。
to_destination_type(x: Value, destination_type: Type) -> Value
是在张量上定义的,它根据type(x)
和destination_type
返回x
的转换值,如下所示:
def to_destination_type(x: Value, destination_type: Type) -> Value:
if type(x) == destination_type:
return x
if is_quantized(destination_type):
if is_quantized(type(x)):
return quantize(x, destination_type)
assert is_float(type(x))
return quantize(x, destination_type)
if is_quantized(type(x)):
assert destination_type = expressed_type(type(x))
return dequantize(type(x))
return convert(x, destination_type)
我们早期讨论了合并 convert
、uniform_quantize
和 uniform_dequantize
操作 (#1576)。合并后,我们不需要上述函数,可以改用 convert
的操作名称。
is_nan(x: Value) -> Value
是在张量上定义的,如果x
的所有元素均为NaN
,则返回true
,否则返回false
。如果x
不是张量,则返回None
。is_sorted(x: Value) -> Value
是在张量上定义的,如果x
元素根据索引的升序字典顺序排序,则返回true
,否则返回false
。如果x
不是张量,则返回None
。is_unique(x: Value) -> Value
在张量上定义,如果x
没有重复元素,则返回true
,否则返回false
。如果x
不是张量,则返回None
。为所有值的所有成员定义
member_name
定义了member_name(x: Value) -> Any
。例如,real_part(x)
会返回相应ComplexConstant
的RealPart
部分。如果x
不是具有适当成员的值,则返回None
。same(x: Value) -> Value
在张量上定义,如果x
的元素彼此均等,则返回true
,否则返回false
。如果张量没有元素,则计为“全部相等”,即函数返回true
。如果x
不是张量,则返回None
。split(x: Value, num_results: Value, axis: Value) -> Value
在张量上定义,并返回沿axis
轴x
的num_results
切片。如果x
不是张量或dim(x, axis) % num_results != 0
,则返回None
。
形状计算
axes(x: Value | Placeholder | Type) -> Value
是range(rank(x))
的快捷方式。dim(x: Value | Placeholder | Type, axis: Value) -> Value
是shape(x)[axis]
的快捷方式。dims(x: Value | Placeholder | Type, axes: List) -> List
是list(map(lambda axis: dim(x, axis), axes))
的快捷方式。index_space(x: Value | Placeholder | Type) -> Value
是在张量上定义的,它返回相应TensorType
的size(x)
索引(按字典顺序升序排列,即[0, ..., 0]
、[0, ..., 1]
、...、shape(x) - 1
)。如果x
不是张量类型、量化张量类型或上述类型之一的值或占位符,则返回None
。rank(x: Value | Placeholder | Type) -> Value
是size(shape(x))
的快捷方式。shape(x: Value | Placeholder | Type) -> Value
通过member_name
在“对类型的函数”部分进行定义。size(x: Value | Placeholder | Type) -> Value
是reduce(lambda x, y: x * y, shape(x))
的快捷方式。
量化计算
def baseline_element_type(x: Value | Placeholder | Type) -> Type
是element_type(baseline_type(x))
的快捷方式。baseline_type
针对张量类型和量化张量类型进行定义,并将它们转换为“基准”,即形状相同但元素类型的量化参数重置为默认值的类型。这可作为一种方便的技巧,用于一致地比较张量和量化张量类型,并且经常需要该方法。对于量化类型,这支持比较忽略量化参数的类型,即shape
、storage_type
、expressed_type
、storage_min
、storage_max
和quantization_dimension
(对于每轴量化类型)必须全部匹配,但scales
和zero points
可能不同。
def baseline_type(x: Value | Placeholder | Type) -> Type:
if type(x) == TensorType:
return x
if type(x) == QuantizedTensorType:
element_type = quantized_tensor_element_type(x)
baseline_element_type = QuantizedTensorElementType(
storage_type = storage_type(element_type),
storage_min = storage_min(element_type),
storage_max = storage_max(element_type),
expressed_type = expressed_type(element_type),
quantization_dimension = quantization_dimension(element_type),
scales = [constant(1.0, expressed_type(element_type))] * dim(x, quantization_dimension(element_type)),
zero_points = [constant(0, storage_type(element_type))] * dim(x, quantization_dimension(element_type)))
return QuantizedTensorType(shape(x), baseline_element_type)
if type(x) is not Type:
return baseline_element_type(type(x))
dequantize
基于量化张量类型定义,并将它们转换为浮点张量类型。方法是使用与量化元素类型关联的零点和标度将表示存储类型的整数值的量化元素转换为表示类型的相应浮点值。
def compute_zero_points(quantized_type, result_type):
if is_per_tensor_quantized(quantized_type):
return broadcast_in_dim(constant(zero_point(quantized_type), storage_type(quantized_type)), [], result_type)
if is_per_axis_quantized(quantized_type):
for i in index_space(result_type):
d = quantization_dimension(quantized_type)
zero_points[i] = zero_points(quantized_type)[i[d]]
return zero_points
def compute_scales(quantized_type, result_type):
if is_per_tensor_quantized(quantized_type):
return broadcast_in_dim(constant(scale(quantized_type), expressed_type(quantized_type)), [],
type(result_type))
if is_per_axis_quantized(quantized_type):
for i in index_space(result_type):
d = quantization_dimension(quantized_type)
scales[i] = scales(quantized_type)[i[d]]
return scales
def dequantize(x: Value) -> Value:
assert is_quantized(x)
x_storage = bitcast_convert(x, storage_type(x))
x_storage_sub = x_storage - compute_zero_points(type(x), type(x_storage))
x_expressed_sub = convert(x_storage_sub, expressed_type(x))
return x_expressed_sub * compute_scales(type(x), type(x_expressed_sub))
quantize
基于浮点张量类型定义,并将它们转换为量化张量类型。此操作通过使用与量化元素类型关联的零点和比例将表示类型的浮点值转换为存储类型的相应整数值来实现。
def quantize(x: Value, type: Type) -> Value:
assert is_float(x) and is_quantized(type)
x_expressed_rounded = round_nearest_even(x / compute_scales(type, type(x)))
x_storage_rounded = convert(x_expressed_rounded, storage_type(type))
x_storage_add = x_storage_rounded + compute_zero_points(type, type(x_storage_rounded))
x_storage = clamp(storage_min(type), x_storage_add, storage_max(type))
return bitcast_convert(x_storage, type)
dequantize_op_quantize
用于指定对量化张量进行元素级计算。对量化元素进行反量化(即将量化元素转换为其表达类型),然后执行操作,再执行量化(即将结果转换回其存储类型)。目前,此函数仅适用于每个张量的量化。正在进行每轴量化 (#1574)。
def dequantize_op_quantize(op, *inputs_and_output_type):
inputs = inputs_and_output_type[:-1]
output_type = inputs_and_output_type[-1]
float_inputs = map(dequantize, inputs)
float_result = op(*float_inputs)
return quantize(float_result, output_type)
def dequantize_batch_norm_grad_or_training_quantize(op, *inputs_and_output_types):
inputs = inputs_and_output_type[:-3]
float_inputs = map(dequantize, inputs)
float_results = op(*float_inputs)
return map(quantize, float_results, inputs_and_output_type[-3:])
def dequantize_compare(lhs, rhs, comparison_direction):
float_lhs = dequantize(lhs)
float_rhs = dequantize(rhs)
return compare(float_lhs, float_rhs, comparison_direction, FLOAT)
def dequantize_select_quantize(pred, on_true, on_false, output_type):
float_on_true = dequantize(on_true)
float_on_false = dequantize(on_false)
float_result = select(pred, float_on_true, float_on_false)
return quantize(float_result, output_type)
网格计算
cross_partition(replica_groups: Value) -> Value
。请参阅上面的“cross_copy”部分。cross_replica(replica_groups: Value) -> Value
。请参阅上面的“cross_copy”部分。cross_replica_and_partition(replica_groups: Value) -> Value
。请参阅上面的“cross_副本_and_partition”部分。flattened_ids(replica_groups: Value) -> Value
。请参阅上面的“flattened_ids”部分。