StableHLO는 머신러닝 (ML) 모델의 상위 수준 작업 (HLO)을 위해 설정된 작업입니다. StableHLO는 다양한 ML 프레임워크와 ML 컴파일러 간의 이동성 레이어로 작동합니다. StableHLO 프로그램을 생성하는 ML 프레임워크는 StableHLO 프로그램을 사용하는 ML 컴파일러와 호환됩니다.
Google의 목표는 다양한 ML 프레임워크 (예: TensorFlow, JAX, PyTorch)와 ML 컴파일러 (예: XLA, IREE) 간에 상호 운용성을 강화하여 ML 개발을 간소화하고 가속화하는 것입니다. 이를 위해 이 문서에서는 StableHLO 프로그래밍 언어의 사양을 제공합니다.
이 사양에는 세 가지 주요 섹션이 있습니다. 먼저 프로그램 섹션에서는 StableHLO 작업으로 구성된 StableHLO 함수로 구성된 StableHLO 프로그램의 구조를 설명합니다. 이 구조 내에서 Ops 섹션은 개별 작업의 시맨틱스를 지정합니다. Execution 섹션은 프로그램 내에서 함께 실행되는 이러한 모든 연산의 의미 체계를 제공합니다. 마지막으로 표기법 섹션에서는 사양 전체에서 사용되는 표기법을 설명합니다.
프로그램
Program ::= {Func}
StableHLO 프로그램은 임의의 수의 StableHLO 함수로 구성됩니다.
다음은 입력 3개(%image
, %weights
, %bias
)와 출력 1개가 있는 @main
함수가 있는 프로그램 예입니다. 함수 본문에는 작업이 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 함수 (이름이 지정된 함수라고도 함)에는 식별자, 입력/출력, 본문이 있습니다. Google은 향후 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개의 차원 (즉, 2개의 축)이 있고 크기는 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)
유형, 값, 양자화 텐서 유형에 0점만 1개 또는 잠재적으로 여러 개일 수 있는지 등 QuantizationZeroPoint
의 시맨틱스에 관한 논의가 진행 중입니다. 이 설명의 결과에 따라 0포인트에 대한 사양은 향후 변경될 수 있습니다. (#1405)
이러한 값과 양자화 텐서 값에 제약 조건을 적용해야 하는지 결정하는 QuantizationStorageMin
및 QuantizationStorageMax
의 시맨틱스에 관한 또 다른 논의가 진행 중입니다. (#1406)
마지막으로, 알 수 없는 크기 크기를 나타내는 방법을 알아볼 계획과 마찬가지로 알 수 없는 배율과 0포인트를 나타내는 방법을 살펴볼 계획입니다. (#1407)
양자화 텐서 유형은 양자화 요소가 있는 텐서를 나타냅니다. 이러한 텐서는 요소에 일반 요소 유형 대신 양자화된 요소 유형이 있다는 점을 제외하면 일반 텐서와 정확히 동일합니다.
양자화 텐서에서 양자화는 텐서별일 수 있습니다. 즉, 전체 텐서에 하나의 scale
및 zero_point
를 보유하거나 축당으로 사용할 수 있습니다. 즉, 특정 차원 quantization_dimension
의 슬라이스당 한 쌍의 scales
와 zero_points
를 보유할 수 있습니다. 좀 더 공식적으로 축당 양자화가 포함된 텐서 t
에는 quantization_dimension
(t[:, ..., 0, ..., :], t[:, ..., 1, ..., :]
등)의 dim(t, quantization_dimension)
슬라이스가 있습니다. 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>>
는 특정 구현에 따라 크게 다를 수 있습니다. 향후 StableHLO에서 튜플 유형을 삭제할 수 있도록 HLO ABI를 변경할 계획입니다. (#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 프로그램은 이러한 유형의 값을 직접 나타낼 수 없습니다. 따라서 T
유형의 스칼라 값을 tensor<T>
유형의 0차원 텐서 값으로 표현하는 것이 관용적입니다.
- 부울 유형은 불리언 값
true
및false
를 나타냅니다. - 정수 유형은 부호 있는 유형 (
si
) 또는 부호 없는 유형 (ui
)일 수 있으며 지원되는 비트 너비 (4
,8
,16
,32
또는64
) 중 하나를 가집니다. 부호 있는siN
유형은-2^(N-1)
이상2^(N-1)-1
까지의 정수 값을 나타내고, 부호 없는uiN
유형은0
부터2^N-1
까지의 정수 값을 나타냅니다. - 부동 소수점 유형은 다음 중 하나입니다.
- 딥 러닝을 위한 FP8 형식에 설명된 FP8 형식의
E4M3
및E5M2
인코딩에 각각 해당하는f8E4M3FN
및f8E5M2
유형입니다. - 심층신경망용 8비트 숫자 형식에 설명된 FP8 형식의
E4M3
및E5M2
인코딩에 해당하는f8E4M3FNUZ
및f8E5M2FNUZ
유형입니다. - 심층신경망을 위한 하이브리드 8비트 부동 소수점 (HFP8) 학습 및 추론에 설명된 FP8 형식의
E4M3
인코딩에 해당하는f8E4M3B11FNUZ
유형입니다. - BFloat16: Cloud TPU의 고성능 비법에 설명된
bfloat16
형식에 해당하는bf16
유형입니다. - IEEE 754 표준에 설명된
binary16
('절반 정밀도'),binary32
('단정밀도') 및binary64
('배정밀도') 형식에 해당하는f16
,f32
,f64
유형입니다.
- 딥 러닝을 위한 FP8 형식에 설명된 FP8 형식의
- 복소수 유형은 동일한 요소 유형의 실수와 허수를 갖는 복합 값을 나타냅니다. 지원되는 복합 유형은
complex<f32>
(두 부분 모두f32
유형임) 및complex<f64>
(두 부분 모두f64
유형임)입니다.
FunctionType ::= '(' InputTypes ')' '->' '(' OutputTypes ')'
InputTypes ::= [ValueType {',' ValueType}]
OutputTypes ::= [ValueType {',' ValueType}]
함수 유형은 이름이 지정된 함수와 익명 함수를 모두 나타냅니다. 여기에는 입력 유형 (->
왼쪽에 있는 유형 목록)과 출력 유형(->
오른쪽에 있는 유형 목록)이 있습니다. 많은 프로그래밍 언어에서 함수 유형은 최고 수준이지만 StableHLO에서는 그렇지 않습니다.
StringType ::= 'string'
문자열 유형은 바이트 시퀀스를 나타냅니다. 많은 프로그래밍 언어와 달리 문자열 유형은 StableHLO의 최고 클래스가 아니며 프로그램 요소의 정적 메타데이터를 지정하는 데만 사용됩니다.
운영
안정적인 HLO 작업 (작업이라고도 함)은 머신러닝 모델에서 상위 수준의 폐쇄형 작업 집합을 나타냅니다. 위에서 설명한 것처럼 StableHLO 구문은 MLIR에서 많은 영감을 받았습니다. 가장 인체공학적인 대안은 아니지만, ML 프레임워크와 ML 컴파일러 간의 상호 운용성을 개선한다는 StableHLO의 목표에 가장 적합합니다.
Op ::= [OpOutputs] OpName OpInputs ':' OpSignature
OpName ::= '"' 'stablehlo' '.' OpMnemonic '"'
OpMnemonic ::= 'abs' | 'add' | ...
StableHLO 작업 (작업이라고도 함)에는 이름, 입력/출력, 서명이 있습니다. 이름은 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 작업의 'Dynamism' 카테고리 - 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) 식별자가 없으므로 (따라서 '익명'이라는 이름) 출력 유형을 선언하지 않습니다 (출력 유형은 함수 내의 return
연산에서 추론됨).
입력 함수의 구문에는 MLIR과의 호환성을 위해 현재 사용되지 않는 부분 (위의 Unused
프로덕션 참고)이 포함됩니다. 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'
정수 상수는 10진수 또는 16진수 표기법을 사용하는 문자열을 통해 정수 값을 나타냅니다. 다른 베이스(예: 바이너리 또는 8진수)는 지원되지 않습니다. 정수 상수에는 다음과 같은 제약 조건이 있습니다.
- (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]
부동 소수점 상수는 십진수 또는 과학 표기법을 사용하는 문자열을 통해 부동 소수점 값을 나타냅니다. 또한 16진수 표기법을 사용하여 기본 비트를 해당 유형의 부동 소수점 형식으로 직접 지정할 수 있습니다. 부동 소수점 상수에는 다음과 같은 제약 조건이 있습니다.
- (C1) 16진수가 아닌 표기법을 사용하는 경우
is_wellformed(float_literal, float_type)
. - (C2) 16진수 표기법을 사용하는 경우
size(hexadecimal_digits) = num_bits(float_type) / 4
.
ComplexConstant ::= ComplexLiteral ':' ComplexType
ComplexLiteral ::= '(' RealPart ',' ImaginaryPart ')'
RealPart ::= FloatLiteral
ImaginaryPart ::= FloatLiteral
복소수 상수는 실수 부분(먼저 나옴)과 허수부 (두 번째로 나옴) 목록을 사용하여 복합 값을 나타냅니다. 예를 들어 (1.0, 0.0) : complex<f32>
는 1.0 + 0.0i
를 나타내고 (0.0, 1.0) : complex<f32>
는 0.0 + 1.0i
를 나타냅니다. 이러한 부분이 메모리에 저장되는 순서는 구현을 통해 정의됩니다. 복잡한 상수에는 다음과 같은 제약 조건이 있습니다.
- (C1)
is_wellformed(real_part, complex_element_type(complex_type))
. - (C2)
is_wellformed(imaginary_part, complex_element_type(complex_type))
.
TensorConstant ::= TensorLiteral ':' TensorType
TensorLiteral ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'
DenseLiteral ::= DenseDimension | DenseElements
DenseDimension ::= '[' [DenseLiteral {',' DenseLiteral}] ']'
DenseElements ::= [ElementLiteral {',' ElementLiteral}]
ElementLiteral ::= BooleanLiteral | IntegerLiteral | FloatLiteral | ComplexLiteral
텐서 상수는 NumPy 표기법을 통해 지정된 중첩 목록을 사용하여 텐서 값을 나타냅니다. 예를 들어 dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
은 색인에서 요소로의 매핑({0, 0} => 1
, {0, 1} => 2
, {0, 2} => 3
, {1, 0} => 4
, {1, 1} => 5
, {1, 2} => 6
)이 있는 텐서 값을 나타냅니다. 이러한 요소가 메모리에 저장되는 순서는 구현으로 정의됩니다. 텐서 상수에는 다음과 같은 제약 조건이 있습니다.
- (C1)
has_syntax(tensor_literal, element_type(tensor_type))
. 각 항목의 의미는 다음과 같습니다.has_syntax(element_literal: Syntax, element_type: Type) = is_wellformed(element_literal, type)
.has_syntax(tensor_literal: List, element_type: Type) = has_syntax(tensor_literal..., element_type)
.
- (C2)
has_shape(tensor_literal, shape(tensor_type))
. 각 항목의 의미는 다음과 같습니다.has_shape(element_literal: Syntax, []) = true
.has_shape(tensor_literal: List, shape: List) = size(tensor_literal) = shape[0] and has_shape(tensor_literal..., shape[1:])
.- 그렇지 않으면
false
입니다.
QuantizedTensorConstant ::= QuantizedTensorLiteral ':' QuantizedTensorType
QuantizedTensorLiteral ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'
양자화 텐서 상수는 요소가 저장소 유형의 상수로 지정된 텐서 상수와 동일한 표기법을 사용하여 양자화 텐서 값을 나타냅니다. 양자화 텐서 상수에는 다음과 같은 제약 조건이 있습니다.
- (C1)
has_syntax(quantized_tensor_literal, storage_type(quantized_tensor_type))
. - (C2)
has_shape(quantized_tensor_literal, shape(quantized_tensor_type))
.
StringConstant ::= StringLiteral
StringLiteral ::= '"' {stringCharacter | escapeSequence} '"'
stringCharacter ::= all ASCII characters except '\00', '\01', ... '\1f' and '"'
escapeSequence ::= '\' ('"' | '\' | 'n' | 't' | (hexadecimalDigit hexadecimalDigit))
문자열 리터럴은 ASCII 문자와 이스케이프 시퀀스를 사용하여 지정된 바이트로 구성됩니다. 인코딩에 구애받지 않으므로 이러한 바이트의 해석은 구현을 통해 정의됩니다. 문자열 리터럴은 string
유형입니다.
작업
abs
시맨틱
operand
텐서에서 요소별 abs 연산을 실행하고 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 유형의 2차원 텐서 상수 |
(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
내에서 다음을 실행합니다.
- 일부 바이너리 트리
schedule
의 경우result@process[result_index] = exec(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 유형 1차원 텐서 상수의 가변 숫자 |
(C1-C3) |
(I3) | channel_id |
si64 유형의 상수 |
(C4) |
(I4) | use_global_device_ids |
i1 유형의 상수 |
(C4) |
(I5) | computation |
함수 | (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
에는is_promotable(element_type(operand), E)
인(tensor<E>, tensor<E>) -> (tensor<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 유형의 2차원 텐서 상수 |
(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 |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (C2), (C4), (C5) |
(I3) | mean |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (C2), (C4) |
(I4) | variance |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (C2), (C4) |
(I5) | grad_output |
부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서 | (C2), (C3) |
(I6) | epsilon |
f32 유형의 상수 |
|
(I7) | feature_index |
si64 유형의 상수 |
(C1), (C5) |
출력
이름 | 유형 | 제약 조건 |
---|---|---|
grad_operand |
부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서 | (C2), (C3) |
grad_scale |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (C2), (C4) |
grad_offset |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (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 |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (C2), (C3) |
(I3) | offset |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (C2), (C4) |
(I4) | mean |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (C5) |
(I5) | variance |
부동 소수점 또는 텐서당 양자화 유형의 1차원 텐서 | (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
차원을 제외한 모든 차원의 평균과 분산을 계산하고 output
, batch_mean
, batch_var
텐서를 생성하는 operand
텐서를 정규화합니다. 보다 공식적으로 이 작업은 다음과 같이 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 |
양자화된 부동 소수점 또는 텐서당 1차원 텐서 | (C2), (C3) |
(I3) | offset |
양자화된 부동 소수점 또는 텐서당 1차원 텐서 | (C2), (C4) |
(I4) | epsilon |
f32 유형의 상수 |
(C1), (C3-C6) |
(I5) | feature_index |
si64 유형의 상수 |
(C1), (C3-C6) |
출력
이름 | 유형 | 제약 조건 |
---|---|---|
output |
부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서 | (C7) |
batch_mean |
양자화된 부동 소수점 또는 텐서당 1차원 텐서 | (C2), (C5) |
batch_var |
양자화된 부동 소수점 또는 텐서당 1차원 텐서 | (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 유형의 1차원 텐서 상수 |
(C2-C6) |
출력
이름 | 유형 | 제약 조건 |
---|---|---|
result |
텐서 또는 양자화 텐서 | (C1), (C3), (C5-C6) |
제약 조건
- (C1)
element_type(result)
는 다음과 같이 주어집니다.!is_per_axis_quantized(operand)
인 경우element_type(operand)
입니다.quantization_dimension(operand)
,scales(operand)
,zero_points(operand)
가quantization_dimension(result)
,scales(result)
,zero_points(result)
응답과 다를 수 있다는 점을 제외하고element_type(operand)
- (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
텐서의 요소별 ceil을 수행하고 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]
콜레스키
시맨틱
행렬 배치의 콜레스키 분해를 계산합니다.
좀 더 공식적으로 index_space(result)
의 모든 i
에 대해 result[i0, ..., iR-3, :, :]
는 하부삼각형(lower
가 true
인 경우) 또는 상삼각형 (lower
가 false
인 경우) 행렬의 형태로 a[i0, ..., iR-3, :, :]
의 콜레스키 분해입니다.
반대쪽 삼각형의 출력 값, 즉 엄격한 위쪽 삼각형 또는 엄격한 아래쪽 삼각형은 구현을 통해 정의됩니다.
입력 행렬이 헤르미티아 양의 정적 행렬이 아닌 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
은 다음과 같이 제공됩니다.
- 프로세스가
process_groups[i]
에 있는i
가 있는 경우operand@process_groups[i, 0]
broadcast_in_dim(constant(0, element_type(result)), [], type(result))
하지 않으면 사용할 수 없습니다.
입력
라벨 | 이름 | 유형 | 제약 조건 |
---|---|---|---|
(I1) | operand |
텐서 | (C3) |
(I2) | replica_groups |
si64 유형 1차원 텐서 상수의 가변 숫자 |
(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 유형의 2차원 텐서 상수 |
(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
또는TOTALORDER
인 경우is_float(element_type(lhs))
입니다.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)
에는E = element_type(lhs)
인complex<E>
유형이 있습니다.
예
// %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
는inputs
의d
번째 차원 크기입니다.
입력
라벨 | 이름 | 유형 | 제약 조건 |
---|---|---|---|
(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
값은 0으로 변환되고 값 true
는 1로 변환됩니다. any-supported-type-to-boolean 변환의 경우 0 값은 false
로 변환되고 0이 아닌 값은 true
로 변환됩니다. 복잡한 유형에서 어떻게 작동하는지는 아래를 참고하세요.
integer-to-integer, integer-to-floating-point 또는 floating-point-to-floating-point와 관련된 변환의 경우 소스 값을 대상 유형에서 정확하게 표현할 수 있는 경우 결과 값은 정확한 표현입니다. 그 외의 경우에는 동작이 미정입니다. (#180)
floating-point-to-integer가 포함된 변환의 경우 소수 부분이 잘립니다. 잘린 값을 대상 유형에 표현할 수 없는 경우 동작은 미정입니다. (#180)
복소수-복소수를 포함하는 변환에서는 실제 부분과 허수부를 변환하는 부동 소수점-부동 소수점 변환과 동일한 동작을 따릅니다.
complex-to-any-other-type 및 any-other-type-to-complex 변환의 경우 소스 허수 값은 각각 무시되거나 대상 허수 값은 각각 0이 됩니다. 실수부의 변환은 부동 소수점 변환을 따릅니다.
원칙적으로 이 작업은 비양자화 (양자화 텐서에서 일반 텐서로의 변환), 양자화 (일반 텐서에서 양자화 텐서로 변환), 재양자화 (양자화 텐서 간 변환)를 표현할 수 있지만, 현재 첫 번째 사용 사례에는 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 유형의 1차원 텐서 상수 |
(C2-C3), (C25) |
(I4) | padding |
si64 유형의 2차원 텐서 상수 |
(C4), (C25) |
(I5) | lhs_dilation |
si64 유형의 1차원 텐서 상수 |
(C5-C6), (C25) |
(I6) | rhs_dilation |
si64 유형의 1차원 텐서 상수 |
(C7-C8), (C25) |
(I7) | window_reversal |
i1 유형의 1차원 텐서 상수 |
(C9) |
(I8) | input_batch_dimension |
si64 유형의 상수 |
(C10), (C13), (C25) |
(I9) | input_feature_dimension |
si64 유형의 상수 |
(C11), (C13-C14) |
(I10) | input_spatial_dimensions |
si64 유형의 1차원 텐서 상수 |
(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 유형의 1차원 텐서 상수 |
(C17-C18), (C25) |
(I14) | output_batch_dimension |
si64 유형의 상수 |
(C20), (C25) |
(I15) | output_feature_dimension |
si64 유형의 상수 |
(C20), (C25), (C33) |
(I16) | output_spatial_dimensions |
si64 유형의 1차원 텐서 상수 |
(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 의 다양한 enum 수입니다. |
(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
텐서의 선행 0비트 수를 요소별로 계산하고 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
시맨틱
inputs
및 called_computations
를 가져와 results
를 생성하는 구현 정의 작업 call_target_name
를 캡슐화합니다. 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 유형의 1차원 텐서 상수 |
(C1), (C3), (C5), (C9), (C12) |
(I4) | rhs_batching_dimensions |
si64 유형의 1차원 텐서 상수 |
(C1), (C4), (C7), (C9) |
(I5) | lhs_contracting_dimensions |
si64 유형의 1차원 텐서 상수 |
(C2), (C3), (C6), (C10) |
(I6) | rhs_contracting_dimensions |
si64 유형의 1차원 텐서 상수 |
(C2), (C4), (C8), (C10) |
(I7) | precision_config |
DEFAULT , HIGH , HIGHEST 의 다양한 enum 수입니다. |
(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
에서 슬라이스를 추출하고 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 |
정수 유형 0차원 텐서의 가변 수 | (C2), (C3) |
(I3) | slice_sizes |
si64 유형의 1차원 텐서 상수 |
(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
시맨틱
start_indices
에서 시작하는 슬라이스가 update
의 값으로 업데이트된다는 점을 제외하고 operand
텐서와 동일한 result
텐서를 생성합니다.
더 공식적으로 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 |
정수 유형 0차원 텐서의 가변 수 | (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
텐서에서 요소별 지수 - 1 연산을 수행하고 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]
fft
시맨틱
실제 및 복잡한 입력/출력에 대해 정방향 및 역 푸리에 변환을 수행합니다.
fft_type
는 다음 중 하나입니다.
FFT
: 복잡에서 복잡하게 FFT를 전달합니다.IFFT
: 복소수에서 복소수로의 역 FFT입니다.RFFT
: 실수에서 복잡한 FFT로 전달합니다.IRFFT
: 실수에서 복소로의 역 FFT입니다 (복소수를 받아 실수를 반환함).
좀 더 공식적으로 복잡한 유형의 1차원 텐서를 입력으로 사용하는 fft
함수를 통해 출력과 동일한 유형의 1차원 텐서를 생성하고 불연속 푸리에 변환을 계산합니다.
fft_type = FFT
의 경우 result
는 L = size(fft_length)
인 일련의 L 계산의 최종 결과로 정의됩니다. 예를 들어 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])
.
또한 동일한 유형 서명을 사용하고 fft
의 역을 계산하는 ifft
함수가 지정된 경우 다음과 같습니다.
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, ..., :])
.
또한 부동 소수점 유형의 1차원 텐서를 사용하는 rfft
함수를 사용하면 동일한 부동 소수점 시맨틱스의 복합 유형의 1차원 텐서를 생성하고 다음과 같이 작동합니다.
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 = size(fft_length)
인 일련의 L 계산의 최종 결과로 정의됩니다. 예를 들어 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])
.
마지막으로, 동일한 유형 서명을 가지고 있고 rfft
의 역을 계산하는 irfft
함수가 지정된 경우 다음과 같습니다.
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 유형의 1차원 텐서 상수 |
(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 유형의 1차원 텐서 상수 |
(C1), (C4-C5), (C13) |
(I4) | collapsed_slice_dims |
si64 유형의 1차원 텐서 상수 |
(C1), (C6-C8), (C13) |
(I5) | start_index_map |
si64 유형의 1차원 텐서 상수 |
(C3), (C9), (C10) |
(I6) | index_vector_dim |
si64 유형의 상수 |
(C2), (C3), (C13) |
(I7) | slice_sizes |
si64 유형의 1차원 텐서 상수 |
(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)
:collapsed_slice_dims
에 해당하는slice_sizes
의 크기 크기가 포함되지 않은 경우를 제외하고combine
은batch_dims
에 해당하는 축에batch_dim_sizes
를 배치하고offset_dims
에 해당하는 축에offset_dim_sizes
를 배치합니다.
- (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 |
함수 | (C1-C3) |
(I3) | false_branch |
함수 | (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
이미지
시맨틱
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]]
아이오타
시맨틱
output
텐서에 iota_dimension
차원을 따라 0부터 오름차순으로 값을 채웁니다. 더 공식적으로
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]]
지도
시맨틱
dimensions
와 함께 inputs
에 맵 함수 computation
를 적용하고 result
텐서를 생성합니다.
더 공식적으로 result[result_index] = computation(inputs...[result_index])
입니다.
dimensions
는 현재 사용되지 않으며 향후 삭제될 수 있습니다. (#487)
입력
라벨 | 이름 | 유형 | 제약 조건 |
---|---|---|---|
(I1) | inputs |
가변 텐서 또는 텐서당 양자화 텐서 | (C1-C4) |
(I2) | dimensions |
si64 유형의 1차원 텐서 상수 |
(C3) |
(I3) | computation |
함수 | (C4) |
출력
이름 | 유형 | 제약 조건 |
---|---|---|
result |
텐서 또는 텐서당 양자화 텐서 | (C1), (C4) |
제약 조건
- (C1)
shape(inputs...) = shape(result)
. - (C2)
0 < size(inputs) = N
. - (C3)
dimensions = range(rank(inputs[0]))
. - (C4)
computation
에는Ei = element_type(inputs[i])
및E' = element_type(result)
인(tensor<E0>, ..., tensor<EN-1>) -> tensor<E'>
유형이 있습니다.
예
// %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
의 요소별 NOT을 실행하고 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
에 종속되는 작업보다 먼저 실행되도록 하고 컴파일러 변환이 배리어를 통해 작업을 이동하지 않도록 합니다. 그 외에는 작업은 ID(예: 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 유형의 1차원 텐서 상수 |
(C1), (C4) |
(I4) | edge_padding_high |
si64 유형의 1차원 텐서 상수 |
(C1), (C4) |
(I5) | interior_padding |
si64 유형의 1차원 텐서 상수 |
(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]
Recv
시맨틱
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
시맨틱
축소 함수 body
를 dimensions
와 함께 inputs
및 init_values
에 적용하고 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:])...)
.- 일부 바이너리 트리
schedule
의 경우reduce(input_slices_converted) = exec(schedule)
이며 각 항목의 의미는 다음과 같습니다.exec(node) = body(exec(node.left), exec(node.right))
.exec(leaf) = leaf.value
.
schedule
는 구현 정의 전체 바이너리 트리이며, 순서 순회는 다음과 같이 구성됩니다.index_space(input_slices_converted)
에 있는 모든index
의 경우input_slices_converted...[index]
값(index
의 사전 오름차순으로 표시됨)- 구현 정의 위치에 구현에서 정의한 양의
init_values_converted
가 산재되어 있습니다.
입력
라벨 | 이름 | 유형 | 제약 조건 |
---|---|---|---|
(I1) | inputs |
가변 텐서 또는 텐서당 양자화 텐서 | (C1-C4), (C6), (C7) |
(I2) | init_values |
0차원 텐서 또는 텐서당 양자화 텐서의 가변 수 | (C2), (C3) |
(I3) | dimensions |
si64 유형의 1차원 텐서 상수 |
(C4), (C5), (C7) |
(I4) | body |
함수 | (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
시맨틱
exponent_bits
및 mantissa_bits
를 사용하고 원래의 부동 소수점 유형으로 다시 operand
를 다른 부동 소수점 유형으로 변환하여 output
텐서를 생성하는 요소별 변환을 수행합니다.
더 공식적으로
- 원래 값의 가수 비트는
roundToIntegralTiesToEven
시맨틱스를 사용하여 원래 값을mantissa_bits
로 나타낼 수 있는 가장 가까운 값으로 반올림하도록 업데이트됩니다. - 그런 다음
mantissa_bits
가 원래 값의 가수 비트 수보다 작으면 가수 비트가mantissa_bits
로 잘립니다. - 그런 다음 중간 결과의 지수 비트가
exponent_bits
에서 제공하는 범위에 맞지 않으면 중간 결과가 원래 부호를 사용하여 무한대로 오버플로되고, 원래 부호를 사용하면 언더플로가 0이 됩니다. - 양자화 유형의 경우
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)
.process_group
의 모든sender
에 대한result@receiver = parts@sender[receiver_index]
. 여기서receiver_index = process_group.index(receiver)
.
입력
라벨 | 이름 | 유형 | 제약 조건 |
---|---|---|---|
(I1) | operand |
텐서 또는 텐서당 양자화 텐서 | (C1), (C2), (C7), (C8) |
(I2) | scatter_dimension |
si64 유형의 상수 |
(C1), (C2), (C8) |
(I3) | replica_groups |
si64 유형의 2차원 텐서 상수 |
(C3-C5) |
(I4) | channel_id |
si64 유형의 상수 |
(C6) |
(I5) | use_global_device_ids |
i1 유형의 상수 |
(C6) |
(I6) | computation |
함수 | (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
에는is_promotable(element_type(operand), E)
인(tensor<E>, tensor<E>) -> (tensor<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 |
0차원 텐서 또는 텐서당 양자화 텐서의 가변 수 | (C1), (C13) |
(I3) | window_dimensions |
si64 유형의 1차원 텐서 상수 |
(C4), (C5), (C15) |
(I4) | window_strides |
si64 유형의 1차원 텐서 상수 |
(C6), (C7), (C15) |
(I5) | base_dilations |
si64 유형의 1차원 텐서 상수 |
(C8), (C9), (C15) |
(I6) | window_dilations |
si64 유형의 1차원 텐서 상수 |
(C10), (C11), (C15) |
(I7) | padding |
si64 유형의 2차원 텐서 상수 |
(C12), (C15) |
(I8) | body |
함수 | (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
에는is_promotable(element_type(inputs[i]), Ei)
가 있는(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ...,
tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)
유형이 있습니다. - (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)
. - 부동 소수점 수: 반올림 속성
roundTowardZero
가 있는 IEEE-754의division(lhs, rhs)
- 복소수는 미정입니다. (#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)
는 다음과 같이 주어집니다.!is_per_axis_quantized(operand)
인 경우element_type(operand)
입니다.quantization_dimension(operand)
와quantization_dimension(result)
가 다를 수 있다는 점을 제외하고element_type(operand)
입니다.
- (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
시맨틱
지정된 dimensions
에 따라 operand
에 있는 요소의 순서를 반대로 바꾸고 result
텐서를 생성합니다. 보다 공식적으로 result[result_index] = operand[operand_index]
각 항목의 의미는 다음과 같습니다.
dimensions
에서d
인 경우operand_index[d] = dim(result, d) - result_index[d] - 1
- 그 밖의 경우에는
operand_index[d] = result_index[d]
입니다.
입력
라벨 | 이름 | 유형 | 제약 조건 |
---|---|---|---|
(I1) | operand |
텐서 또는 텐서당 양자화 텐서 | (C1), (C3) |
(I2) | dimensions |
si64 유형의 1차원 텐서 상수 |
(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]]
RN
시맨틱
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 유형의 1차원 텐서 상수 |
(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 유형의 1차원 텐서 |
(C1), (C2) |
출력
이름 | 유형 | 제약 조건 |
---|---|---|
output_state |
ui64 유형의 1차원 텐서 |
(C1) |
output |
정수 또는 부동 소수점 유형의 텐서 |
제약 조건
- (C1)
type(initial_state) = type(output_state)
. - (C2)
size(initial_state)
는 다음과 같이 정의됩니다.rng_algorithm = DEFAULT
인 경우 구현이 정의됩니다.rng_algorithm = THREE_FRY
인 경우2
입니다.2
또는3
인 경우rng_algorithm = PHILOX
입니다.
예
// %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
시맨틱
scatter_indices
로 지정된 여러 슬라이스가 update_computation
를 사용하여 updates
값으로 업데이트되는 것을 제외하고 inputs
텐서와 동일한 results
텐서를 생성합니다.
다음 다이어그램은 구체적인 예를 사용하여 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...[result_index]
가updated_values...
로 설정된results
의 사본입니다.- 그렇지 않은 경우 다음 단계를 따릅니다.
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 유형의 1차원 텐서 상수 |
(C2), (C4), (C7), (C8) |
(I5) | inserted_window_dims |
si64 유형의 1차원 텐서 상수 |
(C2), (C4), (C9), (C10) |
(I6) | scatter_dims_to_operand_dims |
si64 유형의 1차원 텐서 상수 |
(C11-C13) |
(I7) | index_vector_dim |
si64 유형의 상수 |
(C4), (C11), (C14) |
(I8) | indices_are_sorted |
i1 유형의 상수 |
|
(I9) | unique_indices |
i1 유형의 상수 |
|
(I10) | update_computation |
함수 | (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])
(inserted_window_dims
에 해당하는inputs[0]
의 크기 크기는 포함되지 않음)combine
은update_scatter_dims
에 해당하는 축에update_scatter_dim_sizes
를 배치하고update_window_dims
에 해당하는 축에update_window_dim_sizes
를 배치합니다.
- (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
에는is_promotable(element_type(inputs[i]), Ei)
인(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)
유형이 있습니다. - (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]]
// ]
select
시맨틱
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
의schedule
(reduce 참고)에 init 값이 포함되지 않는다는 점을 제외하고reduce_window
과 정확히 동일하게 작동합니다. 현재 해당 창에 값이 없으면 어떻게 되는지가 지정되어 있지 않습니다. (#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 유형의 1차원 텐서 상수 |
(C2), (C4), (C5) |
(I5) | window_strides |
si64 유형의 1차원 텐서 상수 |
(C2), (C6), (C7) |
(I6) | padding |
si64 유형의 2차원 텐서 상수 |
(C2), (C8) |
(I7) | select |
함수 | (C9) |
(I8) | scatter |
함수 | (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
에는E = element_type(operand)
인(tensor<E>, tensor<E>) -> tensor<i1>
유형이 있습니다. - (C10)
scatter
에는is_promotable(element_type(operand), E)
가 있는(tensor<E>, tensor<E>) -> tensor<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]]
보내기
시맨틱
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
에서 슬라이스를 추출하고 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 유형의 1차원 텐서 상수 |
(C2), (C3), (C5) |
(I3) | limit_indices |
si64 유형의 1차원 텐서 상수 |
(C2), (C3), (C5) |
(I4) | strides |
si64 유형의 1차원 텐서 상수 |
(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
의 1차원 슬라이스를 함께 정렬하고 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
를 반환할 것으로 예상하고 내림차순이 아닌 1차원 슬라이스를 정렬합니다. 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 |
함수 | (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]]
빼기
시맨틱
두 텐서 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]]
Tanh
시맨틱
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 유형의 1차원 텐서 상수 |
(C2-C4) |
출력
이름 | 유형 | 제약 조건 |
---|---|---|
result |
텐서 또는 양자화 텐서 | (C1), (C3-C4) |
제약 조건
- (C1)
element_type(result)
는 다음과 같이 주어집니다.!is_per_axis_quantized(operand)
인 경우element_type(operand)
입니다.quantization_dimension(operand)
와quantization_dimension(result)
가 다를 수 있다는 점을 제외하고element_type(operand)
입니다.
- (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
이 주어지면 result[i0, ..., iR-3, :, :]
는 left_side
가 true
일 때 또는 x * op(a[i0, ..., iR-3, :, :]) = b[i0, ..., iR-3, :, :]
일 때 op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :]
에 대한 해법이며, op(a)
가 transpose_a
에 의해 결정되는 x
변수를 해결합니다. 여기서 op(a)
는 다음 중 하나일 수 있습니다.left_side
false
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
에는Ei = type(val[i])
인tuple<E0, ..., EN-1>
유형이 있습니다.
예
// %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]
동안
시맨틱
cond
함수가 true
를 출력하는 동안 body
함수를 0회 이상 실행하여 출력을 생성합니다. 보다 공식적으로 시맨틱스는 다음과 같이 Python 구문을 사용하여 표현할 수 있습니다.
internal_state = operand
while cond(*internal_state):
internal_state = body(*internal_state)
results = internal_state
무한 루프의 동작은 미정입니다. (#383)
입력
라벨 | 이름 | 유형 | 제약 조건 |
---|---|---|---|
(I1) | operand |
텐서, 양자화 텐서 또는 토큰의 가변 수 | (C1-C3) |
(I2) | cond |
함수 | (C1) |
(I3) | body |
함수 | (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
XOR
시맨틱
두 텐서 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
작업에 루팅된 연산의 그래프를 실행하여 계산됩니다.
실행 순서는 Dataflow와 일치한다면(즉, 작업이 사용 전에 실행되는 경우) 구현으로 정의됩니다. StableHLO에서 모든 부수 효과 작업은 하나의 토큰을 소비하고 하나의 토큰을 생성하므로 (after_all
를 통해 여러 토큰을 하나의 토큰으로 멀티플렉스할 수 있음) 부작용의 실행 순서도 Dataflow와 일치합니다. 위 예시 프로그램의 가능한 실행 순서는 %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
작업을 통해 프로세스 그리드 내의 위치에 액세스할 수 있습니다.
프로세스 그리드 내에서는 프로그램이 모두 같을 수 있고('단일 프로그램, 다중 데이터' 스타일에서) 모두 다를 수도 있고('다중 프로그램, 다중 데이터' 스타일에서) 모두 다를 수도 있습니다. 앞으로 GSPMD를 포함하여 병렬 StableHLO 프로그램을 정의하는 다른 관용구에 대한 지원을 도입할 계획입니다. (#619)
프로세스 그리드 내에서 프로세스는 대부분 서로 독립적입니다. 별도의 작업 상태와 별도의 입력/중간/출력 값을 가지며 아래에 설명된 소수의 집합 작업을 제외하고 대부분의 작업은 프로세스 간에 개별적으로 실행됩니다.
대부분의 작업 실행은 동일한 프로세스의 값만 사용하기 때문에 일반적으로 이러한 값을 이름으로 참조하는 것이 명확합니다.
그러나 집합 작업의 의미 체계를 설명하는 것은 충분하지 않으며 특정 프로세스 내의 name
값을 참조하도록 name@process_id
표기법을 사용하기도 합니다. (이러한 관점에서 정규화되지 않은 name
은 name@(replica_id(), partition_id())
의 약칭으로 볼 수 있습니다.)
프로세스 간 실행 순서는 구현을 통해 정의됩니다. 단, 아래에 설명된 대로 지점 간 통신 및 집합 운영에 의해 도입된 동기화는 예외입니다.
지점 간 통신
StableHLO 프로세스는 StableHLO 채널을 통해 서로 통신할 수 있습니다. 채널은 si64
유형의 양수 ID로 표시됩니다. 다양한 작업을 통해 채널에 값을 전송하고 채널에서 수신할 수 있습니다.
이러한 채널 ID의 출처, 프로세스에서 채널을 인식하는 방법, 프로그램에서 발생하는 동기화 종류와 같은 추가적인 공식화는 미정입니다. (#484)
스트리밍 통신
모든 StableHLO 프로세스는 두 가지 스트리밍 인터페이스에 액세스할 수 있습니다.
- 읽을 수 있는 인피드.
- 작성할 수 있는 아웃피드.
프로세스 간에 통신하여 양쪽 끝에 프로세스가 있는 채널과 달리 인피드 및 아웃피드는 다른 종단 구현이 정의되어 있습니다.
스트리밍 통신이 실행 순서에 미치는 영향 및 이로 인해 도입되는 동기화의 종류와 같은 추가 공식화는 미정입니다. (#484)
집단 운영
StableHLO에는 all_gather
, all_reduce
, all_to_all
, collective_broadcast
, collective_permute
, reduce_scatter
의 6가지 집단 작업이 있습니다. 이러한 모든 작업은 StableHLO 프로세스 그리드의 프로세스를 StableHLO 프로세스 그룹으로 분할하고 다른 프로세스 그룹과 독립적으로 각 프로세스 그룹 내에서 공동 계산을 실행합니다.
각 프로세스 그룹 내에서 집합 작업은 동기화 배리어를 유발할 수 있습니다. 추가 공식화, 예를 들어 이 동기화가 정확히 언제 발생하는지, 프로세스가 이 장벽에 정확히 어떻게 도달하는지, 그렇지 않으면 어떻게 되는지에 관한 자세한 내용은 미정입니다. (#484)
프로세스 그룹에 파티션 간 통신이 포함된 경우(예: 파티션 ID가 다른 프로세스 그룹에 프로세스가 있는 경우) 집합 작업의 실행에는 채널이 필요하며 집합 작업은 si64
유형의 양수 channel_id
를 제공해야 합니다. 교차 복제본 통신에는 채널이 필요하지 않습니다.
집합 오퍼레이션에서 수행되는 계산은 개별 오퍼레이션마다 다르며, 위의 개별 오퍼레이션 섹션에 설명되어 있습니다. 그러나 프로세스 그리드가 프로세스 그룹으로 분할되는 전략은 이러한 작업 간에 공유되며 이 섹션에 설명되어 있습니다. 좀 더 공식적으로 StableHLO는 다음 4가지 전략을 지원합니다.
cross_replica
교차 복제본 통신만 각 프로세스 그룹 내에서 발생합니다. 이 전략은 복제본 ID 목록인 replica_groups
를 사용하고 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
파티션 간 통신만 각 프로세스 그룹 내에서 발생합니다. 이 전략은 파티션 ID 목록인 partition_groups
를 사용하여 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
복제본 간 통신과 파티션 간 통신은 모두 각 프로세스 그룹 내에서 발생할 수 있습니다. 이 전략은 복제본 ID 목록인 replica_groups
를 사용하고 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
이 전략은 replica_id * num_partitions + partition_id
형식의 '평면화된' 프로세스 ID 목록인 flattened_id_groups
를 사용하여 프로세스 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 표준에 의해 정의된 예외 (잘못된 연산, 0으로 나누기, 오버플로, 언더플로 또는 부정확한 예외)를 초래하는 연산은 기본 결과 (표준에 정의된 대로)를 생성하고 상응하는 상태 플래그를 발생시키지 않고 실행을 계속합니다. 표준의 raiseNoFlag
예외 처리와 유사합니다. 비표준 연산 (예: 복잡한 산술 함수, 특정 초월 함수)의 예외는 구현이 정의되어 있습니다.
Notation
구문을 설명하기 위해 이 문서에서는 수정된 ISO 버전의 EBNF 구문 (ISO/IEC 14977:1996, 위키백과)을 사용하며, 두 가지 수정사항이 있습니다. 1) 규칙이 =
가 아닌 ::=
를 사용하여 정의됩니다.
2) 연결은 ,
보다는 병치를 사용하여 표현됩니다.
'유형', '상수', '작업' 섹션 내에서 의미 체계를 설명하기 위해, Google은 아래에 설명된 대로 배열 연산을 간결하게 표현할 수 있도록 지원되는 확장된 Python 문법을 기반으로 하는 수식을 사용합니다. 이 방법은 작은 코드 스니펫에 적합하지만, 드문 경우지만 더 큰 코드 스니펫이 필요한 경우에는 항상 명시적으로 도입되는 바닐라 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으로 변환하는 정확한 형식은 제공하지 않습니다.
일부 특정 수식이 불투명하게 보이는 경우 Google에 알려주시면 이를 개선하도록 노력하겠습니다.
또한 공식에서는 생략 부호를 사용하여 텐서, 텐서 목록 (예: 다양한 수의 텐서에서 발생할 수 있음) 등을 포함한 모든 종류의 목록을 확장합니다. 이는 정확한 공식을 제공하지 않는 또 다른 영역입니다 (예: 목록은 StableHLO 유형 시스템의 일부도 아닌데).
C) 마지막으로 주목할 만한 표기법은 암시적 브로드캐스트입니다. StableHLO opset은 암시적 브로드캐스트를 지원하지 않지만 수식은 간결성을 제공하기 위해 지원합니다. 간단히 말해, 텐서가 예상되는 컨텍스트에서 스칼라가 사용되면 스칼라는 예상 형태로 브로드캐스트됩니다.
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
('유형' 섹션에 정의된 함수)
컨텍스트에 따라 이름은 다른 값을 참조할 수 있습니다. 더 구체적으로 말하면, 연산의 '시맨틱' 섹션 (및 다른 프로그램 요소와 동등한 항목)은 런타임 로직을 정의하므로 모든 입력을 Value
로 사용할 수 있습니다.
반면에 작업 (및 이와 동등한 항목)의 'Constraints' 섹션은 '컴파일 시간' 로직, 즉 일반적으로 런타임 전에 실행되는 로직을 정의하므로 상수 입력만 Value
로 사용할 수 있고 다른 입력은 Placeholder
로만 사용할 수 있습니다.
이름 | '시맨틱'에서 | 'Constraints'에서 |
---|---|---|
전역 함수 | 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
로만 사용할 수 있습니다.
함수
유형 생성
유형을 생성하는 데 사용할 수 있는 함수가 없습니다. 대신 유형 문법이 일반적으로 더 간결하기 때문에 직접 사용합니다. 예를 들어 function_type(
[tensor_type([], E), tensor_type([], E)], [tensor_type([], E)])
가 아닌 (tensor<E>, tensor<E>) -> (tensor<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)
는 텐서 값 2개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(x: Value) -> Any
는 모든 값의 모든 멤버 정의member_name
에 대해 정의됩니다. 예를 들어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
는 양자화 텐서 유형에 정의되며 부동 소수점 텐서 유형으로 변환합니다. 이는 양자화된 요소 유형과 연결된 0점과 배율을 사용하여 저장소 유형의 정수 값을 나타내는 양자화된 요소를 표현된 유형의 상응하는 부동 소수점 값으로 변환함으로써 발생합니다.
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
는 부동 소수점 텐서 유형에 정의되며 양자화 텐서 유형으로 변환합니다. 이는 양자화된 요소 유형과 연결된 0점과 배율을 사용하여 표현된 유형의 부동 소수점 값을 저장소 유형의 상응하는 정수 값으로 변환함으로써 발생합니다.
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_Replica' 섹션을 참조하세요.cross_replica(replica_groups: Value) -> Value
. 위의 'cross_Replica' 섹션을 참조하세요.cross_replica_and_partition(replica_groups: Value) -> Value
. 위의 'cross_Replica_and_partition' 섹션을 참조하세요.flattened_ids(replica_groups: Value) -> Value
. 위의 'flattened_ids' 섹션을 참고하세요.