StableHLO 사양

StableHLO는 머신러닝 (ML) 모델의 고급 작업 (HLO)을 위한 작업 집합입니다. StableHLO는 다양한 ML 프레임워크와 ML 컴파일러 간의 이식성 레이어로 작동합니다. StableHLO 프로그램을 생성하는 ML 프레임워크는 StableHLO 프로그램을 사용하는 ML 컴파일러와 호환됩니다.

Google의 목표는 다양한 ML 프레임워크 (예: TensorFlow, JAX, PyTorch)와 ML 컴파일러 (예: XLA, IREE) 간의 상호 운용성을 높여 ML 개발을 간소화하고 가속화하는 것입니다. 이를 위해 이 문서에서는 StableHLO 프로그래밍 언어의 사양을 제공합니다.

이 사양에는 세 가지 주요 섹션이 포함되어 있습니다. 먼저 프로그램 섹션에서는 StableHLO 프로그램의 구조를 설명합니다. StableHLO 프로그램은 StableHLO 함수로 구성되며 StableHLO 함수는 StableHLO 작업으로 구성됩니다. 이 구조 내에서 Ops 섹션은 개별 작업의 의미를 지정합니다. 실행 섹션에서는 프로그램 내에서 함께 실행되는 이러한 모든 작업의 시맨틱스를 제공합니다. 마지막으로 표기법 섹션에서는 사양 전체에서 사용되는 표기법을 설명합니다.

이전 버전의 StableHLO 사양을 보려면 원하는 태그가 지정된 버전에서 저장소를 여세요. 예를 들어 StableHLO v0.19.0 사양을 참고하세요. StableHLO의 각 부버전 업데이트에서 발생한 변경사항을 확인하려면 VhloDialect.td의 버전 로그를 참고하세요.

프로그램

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 함수 (이름이 지정된 함수라고도 함)에는 식별자, 입력/출력, 본문이 있습니다. 향후 함수에 추가 메타데이터를 도입하여 HLO와의 호환성을 개선할 계획입니다 (#425, #626, #740, #744).

식별자

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

StableHLO 식별자는 여러 프로그래밍 언어의 식별자와 유사하며 두 가지 특이사항이 있습니다. 1) 모든 식별자에는 다양한 종류의 식별자를 구분하는 시길이 있습니다. 2) 값 식별자는 StableHLO 프로그램 생성을 단순화하기 위해 완전히 숫자일 수 있습니다.

유형

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

StableHLO 유형은 StableHLO 값을 나타내는 값 유형 (일급 유형이라고도 함)과 기타 프로그램 요소를 설명하는 비값 유형으로 분류됩니다. StableHLO 유형은 여러 프로그래밍 언어의 유형과 유사하며, 주요 특이점은 StableHLO의 도메인별 특성으로 인해 일부 비정상적인 결과가 발생한다는 것입니다 (예: 스칼라 유형은 값 유형이 아님).

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

텐서 유형은 텐서, 즉 다차원 배열을 나타냅니다. 모양과 요소 유형이 있으며, 여기서 모양은 0부터 R-1까지 번호가 매겨진 해당 차원 (축이라고도 함)의 오름차순으로 음수가 아닌 차원 크기 또는 알 수 없는 차원 크기를 나타냅니다. 차원 수 R을 순위라고 합니다. 예를 들어 tensor<2x3xf32>은 모양이 2x3이고 요소 유형이 f32인 텐서 유형입니다. 크기가 2와 3인 두 개의 차원(즉, 두 개의 축)인 0번째 차원과 1번째 차원이 있습니다. 순위는 2입니다.

모양은 부분적으로 또는 완전히 알 수 없을 수 있습니다 (동적). 예를 들어 tensor<?x2xf64>는 부분적으로 알 수 없고 tensor<?x?xf64>는 완전히 알 수 없습니다. 동적 측정기준 크기는 ?를 사용하여 표현됩니다. 모양은 순위를 해제할 수 없습니다.

향후에는 텐서 유형을 차원 크기 및 요소 유형을 넘어 레이아웃(#629) 및 희소성(#1078)을 포함하도록 확장할 계획입니다.

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

양자화된 요소 유형은 표현된 유형의 부동 소수점 값에 해당하는 storage_min~storage_max (포함) 범위의 저장 유형 정수 값을 나타냅니다. 주어진 정수 값 i의 경우 해당 부동 소수점 값 f은 f = (i - zero_point) * scale로 계산할 수 있습니다. 여기서 scale 및 zero_point을 양자화 매개변수라고 합니다. storage_min 및 storage_max는 문법에서 선택사항이지만 기본값은 각각 min_value(storage_type) 및 max_value(storage_type)입니다. 양자화된 요소 유형에는 다음과 같은 제약 조건이 있습니다.

• (C1) type(storage_min) = storage_type.
• (C2) type(storage_max) = storage_type
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type).
• (C4) type(scales...) = expressed_type
• (C5) 0 < scales.
• (C6) is_finite(scales...)
• (C7) storage_min <= zero_points <= storage_max.
• (C8) type(zero_points...) = storage_type
• (C9) size(scales) = size(zero_points).
• (C10) is_empty(quantization_dimension)인 경우 size(scales) = 1입니다.
• (C11) 0 <= quantization_dimension.

현재 QuantizationScale는 부동 소수점 상수이지만, 승수와 시프트로 표현되는 정수 기반 스케일에 대한 관심이 높습니다. Google에서는 가까운 시일 내에 이 문제를 살펴볼 계획입니다(#1404).

양자화된 텐서 유형의 유형, 값, 0점이 하나만 있을 수 있는지 아니면 여러 개 있을 수 있는지 등 QuantizationZeroPoint의 의미에 관한 논의가 진행 중입니다. 이 토론의 결과에 따라 0점 관련 사양이 향후 변경될 수 있습니다 (#1405).

또 다른 진행 중인 논의는 QuantizationStorageMin 및 QuantizationStorageMax의 시맨틱을 포함하여 이러한 값과 양자화된 텐서의 값에 제약 조건을 적용해야 하는지 결정하는 것입니다(#1406).

마지막으로 알 수 없는 스케일과 0점을 표현하는 방법을 알 수 없는 차원 크기를 표현하는 방법과 유사하게 살펴볼 계획입니다 (#1407).

양자화된 텐서 유형은 양자화된 요소가 있는 텐서를 나타냅니다. 이러한 텐서는 일반 텐서와 정확히 동일하지만 요소가 일반 요소 유형 대신 양자화된 요소 유형을 갖습니다.

양자화된 텐서에서 양자화는 텐서별일 수 있습니다. 즉, 전체 텐서에 하나의 scale 및 zero_point가 있거나 축별일 수 있습니다. 즉, 특정 차원 quantization_dimension의 슬라이스당 한 쌍의 여러 scales 및 zero_points가 있습니다. 더 공식적으로 축별 양자화가 적용된 텐서 t에는 quantization_dimension의 슬라이스가 dim(t, quantization_dimension)개 있습니다(t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] 등). i번째 슬라이스의 모든 요소는 scales[i] 및 zero_points[i]를 양자화 매개변수로 사용합니다. 양자화된 텐서 유형에는 다음과 같은 제약이 있습니다.

• 텐서별 양자화의 경우 다음을 충족해야 합니다.
  • 추가 제약 조건이 없습니다.
• 축별 양자화의 경우 다음을 충족해야 합니다.
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

토큰 유형은 토큰, 즉 일부 작업에서 생성하고 사용하는 불투명한 값을 나타냅니다. 토큰은 실행 섹션에 설명된 대로 작업에 실행 순서를 적용하는 데 사용됩니다.

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

버퍼 유형은 버퍼를 나타냅니다. 예를 들어 XLA에서 버퍼는 일관된 스토리지가 있는 다차원 배열입니다. 텐서 유형과 마찬가지로 버퍼 유형에는 모양과 요소 유형이 있습니다. 여기서 모양은 0부터 R-1까지 번호가 매겨진 해당 차원 (축이라고도 함)의 오름차순으로 음수가 아닌 차원 크기 또는 알 수 없는 차원 크기를 나타냅니다. 측정기준 수 R를 순위라고 합니다. 예를 들어 memref<2x3xf32>은 모양이 2x3이고 요소 유형이 f32인 버퍼 유형입니다. 크기가 2와 3인 두 개의 차원 (즉, 두 개의 축)인 0번째 차원과 1번째 차원이 있습니다. 순위는 2입니다.

버퍼는 custom_call~CreateBuffer 또는 Pin를 사용하여 할당하고 custom_call~Unpin를 통해 할당 해제할 수 있습니다. custom_call 작업만 버퍼 내부의 콘텐츠를 읽고 쓸 수 있습니다. 자세한 내용은 custom_call을 참고하세요.

튜플 유형은 튜플, 즉 이질적인 목록을 나타냅니다. 튜플은 HLO와의 호환성을 위해서만 존재하는 기존 기능입니다. HLO에서 튜플은 가변 입력과 출력을 나타내는 데 사용됩니다. StableHLO에서는 가변 입력과 출력이 기본적으로 지원되며 StableHLO에서 튜플을 사용하는 유일한 경우는 특정 구현에 따라 T, tuple<T>, tuple<tuple<T>>이 실질적으로 다를 수 있는 HLO ABI를 포괄적으로 표현하는 경우입니다. 향후 HLO ABI를 변경하여 StableHLO에서 튜플 유형을 삭제할 수 있도록 할 계획입니다(#598).

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

요소 유형은 텐서 유형의 요소를 나타냅니다. 많은 프로그래밍 언어와 달리 이러한 유형은 StableHLO에서 최상위 클래스가 아닙니다. 즉, StableHLO 프로그램은 이러한 유형의 값을 직접 나타낼 수 없습니다. 따라서 T 유형의 스칼라 값을 tensor<T> 유형의 0차원 텐서 값으로 나타내는 것이 관용적입니다.

• 불리언 유형은 불리언 값 true 및 false를 나타냅니다.
• 정수 유형은 부호가 있거나 (si) 부호가 없을 수 있으며 (ui) 지원되는 비트 너비 (2, 4, 8, 16, 32 또는 64) 중 하나를 가집니다. 부호 있는 siN 유형은 -2^(N-1)부터 2^(N-1)-1까지의 정수 값을 나타내고 부호 없는 uiN 유형은 0부터 2^N-1까지의 정수 값을 나타냅니다.
• 부동 소수점 유형은 다음 중 하나일 수 있습니다.
  • IEEE-754 규칙을 따르는 f8E3M4, f8E4M3, f8E5M2 8비트 부동 소수점 숫자입니다.
  • 딥 러닝을 위한 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 유형입니다.
  • tf32 유형은 TensorFloat32 형식에 해당하며 StableHLO에서 지원이 제한됩니다.
  • OCP 마이크로스케일링 형식 사양에 설명된 f4E2M1FN, f6E2M3FN, f6E3M2FN, f8E8M0FNU MX (마이크로스케일링) 유형
• 복합 유형은 동일한 요소 유형의 실수부와 허수부가 있는 복합 값을 나타냅니다. 지원되는 복잡한 유형은 complex<f32> (두 부분이 모두 f32 유형임) 및 complex<f64>(두 부분이 모두 f64 유형임)입니다.

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

함수 유형은 이름이 지정된 함수와 익명 함수를 모두 나타냅니다. 입력 유형 (->의 왼쪽에 있는 유형 목록)과 출력 유형(->의 오른쪽에 있는 유형 목록)이 있습니다. 많은 프로그래밍 언어에서 함수 유형은 일급이지만 StableHLO에서는 그렇지 않습니다.

StringType ::= 'string'

문자열 유형은 바이트 시퀀스를 나타냅니다. 많은 프로그래밍 언어와 달리 문자열 유형은 StableHLO에서 일급이 아니며 프로그램 요소의 정적 메타데이터를 지정하는 데만 사용됩니다.

작업

StableHLO 작업 (작업이라고도 함)은 머신러닝 모델의 폐쇄된 고급 작업 집합을 나타냅니다. 위에서 설명한 것처럼 StableHLO 구문은 MLIR에서 많은 영감을 받았습니다. MLIR이 반드시 가장 인체공학적인 대안은 아니지만 ML 프레임워크와 ML 컴파일러 간의 상호 운용성을 높이려는 StableHLO의 목표에 가장 적합하다고 할 수 있습니다.

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

StableHLO 작업 (작업이라고도 함)에는 이름, 입력/출력, 서명이 있습니다. 이름은 stablehlo. 접두사와 지원되는 작업을 고유하게 식별하는 니모닉으로 구성됩니다. 지원되는 모든 작업의 전체 목록은 아래를 참고하세요.

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

작업은 입력을 사용하고 출력을 생성합니다. 입력은 입력 값 (실행 중에 계산됨), 입력 함수 (StableHLO에서는 함수가 최고 수준 값이 아니므로 정적으로 제공됨), 입력 속성 (정적으로 제공됨)으로 분류됩니다. 작업에서 소비하고 생성하는 입력 및 출력 종류는 니모닉에 따라 다릅니다. 예를 들어 add 작업은 입력 값 2개를 사용하고 출력 값 1개를 생성합니다. 이에 비해 select_and_scatter 작업은 입력 값 3개, 입력 함수 2개, 입력 속성 3개를 사용합니다.

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

입력 함수 (익명 함수라고도 함)는 다음을 제외하고 이름이 지정된 함수와 매우 유사합니다. 1) 식별자가 없습니다('익명'이라는 이름이 붙은 이유). 2) 출력 유형을 선언하지 않습니다 (출력 유형은 함수 내의 return 작업에서 추론됨).

입력 함수의 구문에는 현재 사용되지 않는 부분이 포함되어 있습니다 (위의 Unused 프로덕션 참고). 이는 MLIR과의 호환성을 위한 것입니다. MLIR에는 점프 작업으로 연결된 여러 작업 '블록'이 있을 수 있는 '리전'이라는 더 일반적인 개념이 있습니다. 이러한 블록에는 Unused 프로덕션에 해당하는 ID가 있어 서로 구분할 수 있습니다. StableHLO에는 점프 작업이 없으므로 MLIR 구문의 해당 부분은 사용되지 않습니다 (하지만 여전히 존재함).

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

입력 속성에는 이름과 지원되는 상수 중 하나의 값이 있습니다. 이는 프로그램 요소의 정적 메타데이터를 지정하는 기본 방법입니다. 예를 들어 concatenate 작업은 dimension 속성을 사용하여 입력 값이 연결되는 차원을 지정합니다. 마찬가지로 slice 작업은 start_indices, limit_indices과 같은 여러 속성을 사용하여 입력 값을 슬라이스하는 데 사용되는 경계를 지정합니다.

현재 실제 StableHLO 프로그램에는 이 문서에 설명되지 않은 속성이 포함되는 경우가 있습니다. 향후 이러한 속성을 StableHLO opset에 흡수하거나 StableHLO 프로그램에 표시되지 않도록 금지할 계획입니다. 그동안 다음 속성 목록을 확인하세요.

• layout (#629)
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740)
• 위치 메타데이터 (#594)

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

작업 서명은 모든 입력 값의 유형 (->의 왼쪽 유형 목록)과 모든 출력 값의 유형 (->의 오른쪽 유형 목록)으로 구성됩니다. 엄밀히 말하면 입력 유형은 중복되고 출력 유형도 거의 항상 중복됩니다 (대부분의 StableHLO 작업의 경우 출력 유형은 입력에서 추론할 수 있기 때문). 하지만 MLIR과의 호환성을 위해 op 서명은 의도적으로 StableHLO 구문의 일부입니다.

아래는 니모닉이 select_and_scatter인 작업의 예입니다. 3개의 입력 값 (%operand, %source, %init_value), 2개의 입력 함수, 3개의 입력 속성 (window_dimensions, window_strides, padding)을 사용합니다. 작업의 서명에는 입력 값의 유형만 포함됩니다(인라인으로 제공되는 입력 함수 및 속성의 유형은 포함되지 않음).

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

상수

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

StableHLO 상수에는 리터럴과 유형이 있으며, 이 둘이 함께 StableHLO 값을 나타냅니다. 일반적으로 유형은 상수 구문의 일부입니다. 단, 명확한 경우 (예: 불리언 상수는 명확하게 i1 유형을 갖는 반면 정수 상수는 여러 가능한 유형을 가질 수 있음)는 예외입니다.

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

불리언 상수는 불리언 값 true 및 false을 나타냅니다. 불리언 상수의 유형은 i1입니다.

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

정수 상수는 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))

입력

출력

제약조건

• (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)< : (t>ens>or3xi32<) - t>ensor3xi32
// %result: [2, 0, 2]

 예시 더보기

추가

시맨틱스

두 텐서 lhs와 rhs의 요소별 덧셈을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 불리언: 논리합
• 정수의 경우 정수 덧셈입니다.
• 부동 소수점의 경우: IEEE-754의 addition
• 복소수의 경우 복소수 덧셈입니다.
• 양자화된 유형의 경우: dequantize_op_quantize(add, lhs, rhs, type(result))

입력

출력

제약조건

• 작업에서 양자화되지 않은 텐서를 사용하는 경우 다음을 충족해야 합니다.
  • (C1) type(lhs) = type(rhs) = type(result).
• 작업에서 양자화된 텐서를 사용하는 경우 다음을 충족해야 합니다.
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result)
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result).
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result).
  • (C6) is_per_axis_quantized(lhs) 시 quantization_dimension(lhs) = quantization_dimension(result)
  • (C7) is_per_axis_quantized(rhs) 시 quantization_dimension(rhs) = quantization_dimension(result).

예

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

 예시 더보기

after_all

시맨틱스

inputs를 생성하는 작업이 result에 종속된 작업보다 먼저 실행되도록 합니다. 이 작업의 실행은 아무것도 하지 않으며 result에서 inputs로의 데이터 종속 항목을 설정하기 위해서만 존재합니다.

입력

출력

예

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

 예시 더보기

all_gather

시맨틱스

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 각 프로세스의 operands 텐서 값을 all_gather_dim에 따라 연결하고 results 텐서를 생성합니다.

이 작업은 StableHLO 프로세스 그리드를 다음과 같이 정의된 process_groups로 분할합니다.

• 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 = false인 경우
• flattened_ids(replica_groups) channel_id > 0 and use_global_device_ids = true인 경우

그런 다음 각 process_group 내에서 다음을 실행합니다.

• process_group의 모든 receiver에 대해 operands...@receiver = [operand@sender for sender in process_group]
• process_group의 모든 process에 대해 results...@process = concatenate(operands...@process, all_gather_dim)

입력

출력

제약조건

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups)
• (C3) size(replica_groups)는 다음과 같이 정의됩니다.
  • cross_replica이 사용된 경우 num_replicas
  • cross_replica_and_partition이 사용된 경우 num_replicas
  • flattened_ids이 사용된 경우 num_processes
• (C4) 0 <= replica_groups < size(replica_groups)
• (C5) use_global_device_ids = true 시 channel_id > 0.
• (C6) type(results...) = type(operands...)(다음 제외)
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

예

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [[1, 2], [3, 4]]
// %operand0@(1, 0): [[5, 6], [7, 8]]
// %operand1@(0, 0): [[11, 12], [13, 14]]
// %operand1@(1, 0): [[15, 16], [17, 18]]
%result:2 = "stablehlo.all_gather"(%operand0, %operand1) {
  all_gather_dim = 1 : i64,
  replica_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64,
  // channel_id = 0
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
  // use_global_device_ids = false
}< : (ten>sor2x2xi<64, ten>sor>2x2xi64)< - (ten>sor2x4xi<64, ten>sor2x4xi64)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

 예시 더보기

all_reduce

시맨틱스

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 각 프로세스의 operands 텐서 값에 감소 함수 computation를 적용하고 results 텐서를 생성합니다.

이 작업은 StableHLO 프로세스 그리드를 다음과 같이 정의된 process_groups로 분할합니다.

• 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 = false인 경우
• flattened_ids(replica_groups) channel_id > 0 and use_global_device_ids = true인 경우

그런 다음 각 process_group 내에서 다음을 실행합니다.

• results...@process[result_index] = exec(schedule)(일부 이진 트리 schedule의 경우) 여기서
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule는 중위 순회 순서가 to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0]))인 구현 정의 바이너리 트리입니다.

입력

출력

제약조건

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups)는 다음과 같이 정의됩니다.
  • cross_replica이 사용된 경우 num_replicas
  • cross_replica_and_partition이 사용된 경우 num_replicas
  • flattened_ids이 사용된 경우 num_processes
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) use_global_device_ids = true 시 channel_id > 0.
• (C5) computation의 유형은 (tensor<E>, tensor<E>) -> (tensor<E>)이며 여기서 is_promotable(element_type(operand), E)입니다.
• (C6) shape(results...) = shape(operands...)
• (C7) element_type(results...) = E.

예

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [1, 2, 3, 4]
// %operand0@(1, 0): [5, 6, 7, 8]
// %operand1@(0, 0): [9, 10, 11, 12]
// %operand1@(1, 0): [13, 14, 15, 16]
%result:2 = "stablehlo.all_reduce"(%operand0, %operand0) ({
  ^bb0(%ar<g0:> tensori64, %ar<g1:> tensori64):
    %0 = "stablehlo.add"(%ar<g0,> %arg1) <: (>ten>sori64,< te>nsori64) - tensori64
    "stable<hlo>.re>turn"(%0) : (tensori64) - ()<
}) {
  >replica_g<roups => dense[[0, 1]] : tensor1x2xi64,
  // channel_id = 0
  channel_hand<le = #stablehlo.chan>nel_handlehandle = 0, type = 0
  // use_global_<devic>e_ids = <false>
} >: (tenso<r4xi6>4, tenso<r4xi6>4) - (tensor4xi64, tensor4xi64)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

 예시 더보기

all_to_all

시맨틱스

StableHLO 프로세스 그리드의 각 프로세스 그룹 내에서 operands 텐서의 값을 split_dimension을 따라 부분으로 분할하고, 분할된 부분을 프로세스 간에 분산하고, 분산된 부분을 concat_dimension을 따라 연결하여 results 텐서를 생성합니다. 이 작업은 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(operands...@sender, split_count, split_dimension)
• receiver_index = process_group.index(receiver)인 경우 scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group]
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

입력

출력

제약조건

• (C1) 0 <= split_dimension < rank(operands...).
• (C2) dim(operands..., split_dimension) % split_count = 0
• (C3) 0 <= concat_dimension < rank(operands...).
• (C4) 0 < split_count
• (C5) is_unique(replica_groups).
• (C6) size(replica_groups)는 다음과 같이 정의됩니다.
  • cross_replica이 사용된 경우 num_replicas
  • cross_partition이 사용된 경우 num_partitions
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count
• (C9) split_dimension != concat_dimension인 경우를 제외하고 type(results...) = type(operands...):
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

예

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

 예시 더보기

및

시맨틱스

두 텐서 lhs와 rhs의 요소별 AND를 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 불리언: 논리곱
• 정수의 경우 비트별 AND입니다.

입력

출력

제약조건

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

예

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

 예시 더보기

atan2

시맨틱스

lhs 및 rhs 텐서에 요소별 atan2 연산을 수행하여 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점의 경우: IEEE-754의 atan2
• 복소수의 경우 복소수 atan2입니다.
• 양자화된 유형의 경우: dequantize_op_quantize(atan2, lhs, rhs, type(result))

입력

출력

제약조건

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

예

// %lhs: [0.0, 1.0, -1.0]
// %rhs: [0.0, 0.0, 0.0]
%result = "stablehlo.atan2"(%lhs, %rhs)< : (t>ensor3xf<64, t>ens>or3xf64<) - t>ensor3xf64
// %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))를 실행합니다.

입력

출력

제약조건

• (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
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf64,
 <    tenso>r2x>2x2xf64)< - (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %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))를 실행합니다.

입력

출력

제약조건

• (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
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ens>or2xf64<) - tenso>r2x2x2xf64
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

batch_norm_training

시맨틱스

feature_index 차원을 제외한 모든 차원에서 평균과 분산을 계산하고 operand 텐서를 정규화하여 output, batch_mean, batch_var 텐서를 생성합니다. 더 공식적으로 이 작업은 다음과 같이 Python 구문을 사용하여 기존 StableHLO 작업으로 분해할 수 있습니다.

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

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

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

양자화된 유형의 경우 dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var))를 실행합니다.

입력

출력

제약조건

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var, output의 baseline_element_type이 동일합니다.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index)
• (C5) size(batch_mean) = dim(operand, feature_index).
• (C6) size(batch_var) = dim(operand, feature_index)
• (C7) baseline_type(output) = baseline_type(operand).

예

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
%output, %batch_mean, %batch_var = "stablehlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ens>or2xf64) -
 <   (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %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 텐서에서 비트캐스트 작업을 실행하고 전체 operand 텐서의 비트가 result 텐서의 유형을 사용하여 재해석되는 result 텐서를 생성합니다.

더 공식적으로 E = element_type(operand), E' = element_type(result), R = rank(operand)가 주어졌을 때:

• num_bits(E') < num_bits(E)인 경우 bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1])입니다.
• num_bits(E') > num_bits(E)인 경우 bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :])입니다.
• num_bits(E') = num_bits(E)인 경우 bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1])입니다.

bits는 지정된 값의 메모리 내 표현을 반환하며 텐서의 정확한 표현이 구현에 따라 정의되고 요소 유형의 정확한 표현도 구현에 따라 정의되므로 동작이 구현에 따라 정의됩니다.

입력

출력

제약조건

• (C1) E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result), R = rank(operand)가 주어졌을 때:
  • num_bits(E') = num_bits(E)인 경우 shape(result) = shape(operand)
  • num_bits(E') < num_bits(E)인 경우:
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i)(모든 0 <= i < R)
  • dim(result, R) * num_bits(E') = num_bits(E).
  • num_bits(E') > num_bits(E)인 경우:
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i)(모든 0 <= i < R)
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) is_complex(operand) or is_complex(result)인 경우 is_complex(operand) and is_complex(result)입니다.

예

// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand)< : >(te>nsorf64<) - t>ensor4xf16
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

 예시 더보기

broadcast_in_dim

시맨틱스

operand 텐서의 데이터를 복제하여 입력 텐서의 차원 또는 순위를 확장하고 result 텐서를 생성합니다. 더 공식적으로 말하면 result[result_index] = operand[operand_index]입니다. 여기서 axes(operand)의 모든 d에 대해 다음이 적용됩니다.

• dim(operand, d) = 1인 경우 operand_index[d] = 0
• 그 밖의 경우에는 operand_index[d] = result_index[broadcast_dimensions[d]]입니다.

입력

출력

제약조건

• (C1) element_type(result)은 다음 식으로 주어집니다.
  • !is_per_axis_quantized(operand)인 경우 element_type(operand)
  • element_type(operand)(quantization_dimension(operand), scales(operand), zero_points(operand)이 각각 quantization_dimension(result), scales(result), zero_points(result)와 다를 수 있음)
• (C2) size(broadcast_dimensions) = rank(operand)
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions)
• (C5) axes(operand)의 모든 d에 대해 다음을 실행합니다.
  • dim(operand, d) = 1 또는
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • dim(operand, quantization_dimension(operand)) = 1이면 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))입니다.

예

// %operand: [
//            [1, 2, 3]
//           ]
%result = "stablehlo.broadcast_in_dim"(%operand) {
  broadcast_dimensio<ns = arra>yi64: 2, 1
}< : (ten>sor>1x3xi32<) - tenso>r2x3x2xi32
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 예시 더보기

케이스

시맨틱스

index 값에 따라 branches에서 정확히 하나의 함수를 실행한 결과를 생성합니다. 더 격식 있게 표현하면 다음과 같습니다. result = selected_branch() 여기서 각 항목의 의미는 다음과 같습니다.

• 0 <= index < size(branches)인 경우 selected_branch = branches[index]
• 그 밖의 경우에는 selected_branch = branches[-1]입니다.

입력

출력

제약조건

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

예

// %index: -1
// %result_branch0: [0, 0]
// %result_branch1: [1, 1]
%result0, %result1 = "stablehlo.case"(%index) ({
  "stablehlo.return"(%result_branch0, %resul<t_bra>nch0) : <(tens>or2>xi64, tensor2xi64) - ()
}, {
  "stablehlo.return"(%result_branc<h1, %>result_b<ranch>1) >: (tensor2xi64, <ten>sor>2xi64) -< ()
}>) : (ten<sori3>2) - (tensor2xi64, tensor2xi64)
// %result0: [1, 1]
// %result1: [1, 1]

 예시 더보기

cbrt

시맨틱스

operand 텐서에 대해 요소별 세제곱근 연산을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점의 경우: IEEE-754의 rootn(x, 3)
• 복소수의 경우 복소수 세제곱근입니다.
• 양자화된 유형: dequantize_op_quantize(cbrt, operand, type(result))

입력

출력

제약조건

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

예

// %operand: [0.0, 1.0, 8.0, 27.0]
%result = "stablehlo.cbrt"(%operand)< : (t>ens>or4xf64<) - t>ensor4xf64
// %result: [0.0, 1.0, 2.0, 3.0]

 예시 더보기

ceil

시맨틱스

operand 텐서의 요소별 올림을 실행하고 result 텐서를 생성합니다. IEEE-754 사양의 roundToIntegralTowardPositive 작업을 구현합니다. 양자화된 유형의 경우 dequantize_op_quantize(ceil, operand, type(result))를 실행합니다.

입력

출력

제약조건

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

예

// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.ceil"(%operand)< : (t>ens>or5xf32<) - t>ensor5xf32
// %result: [-0.0, -0.0, 1.0, 1.0, 2.0]

 예시 더보기

cholesky

시맨틱스

행렬 배치에 대한 Cholesky 분해를 계산합니다.

더 공식적으로 말하면 index_space(result)의 모든 i에 대해 result[i0, ..., iR-3, :, :]는 a[i0, ..., iR-3, :, :]의 촐레스키 분해이며, 하삼각(lower이 true인 경우) 또는 상삼각 (lower이 false인 경우) 행렬 형식입니다. 반대쪽 삼각형, 즉 엄격한 상부 삼각형 또는 엄격한 하부 삼각형의 출력 값은 구현에 따라 정의됩니다.

입력 행렬이 에르미트 양의 정부호 행렬이 아닌 i이 있으면 동작이 정의되지 않습니다.

양자화된 유형의 경우 dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result))를 실행합니다.

입력

출력

제약조건

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

예

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

고정하다

시맨틱스

operand 텐서의 모든 요소를 최소값과 최대값 사이에 고정하고 result 텐서를 생성합니다. 더 공식적으로는 result[result_index] = minimum(maximum(operand[result_index], min_element), max_element)입니다. 여기서 min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]입니다. 양자화된 유형의 경우 dequantize_op_quantize(clamp, min, operand, max, type(result))을 실행합니다.

복소수에 순서를 적용하면 놀라운 의미가 발생하므로 향후 이 작업에 대한 복소수 지원을 삭제할 계획입니다 (#560).

입력

출력

제약조건

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

예

// %min: [5, 10, 15]
// %operand: [3, 13, 23]
// %max: [10, 15, 20]
%result = "stablehlo.clamp"(%min, %operand, %max)< : (t>ensor3xi<32, t>ensor3xi<32, t>ens>or3xi32<) - t>ensor3xi32
// %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(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result))을 반환합니다.

입력

출력

제약조건

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N(여기서 N은 다음과 같이 정의됨)
  • cross_replica이 사용된 경우 num_replicas
  • cross_partition이 사용된 경우 num_partitions
• (C3) type(result) = type(operand).

예

// num_replicas: 4
// num_partitions: 1
// %operand@(0, 0): [[1, 2]]
// %operand@(1, 0): [[3, 4]]
// %operand@(2, 0): [[5, 6]]
// %operand@(3, 0): [[7, 8]]
%result = "stablehlo.collective_broadcast"(%operand) {
  replica_grou<ps = den>se[[2, 1]<] : ten>sor1x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
} : (ten>sor>1x2xi64<) - ten>sor1x2xi64
// %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))을 반환합니다.

입력

출력

제약조건

• (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_pai<rs = dense[[0, 1>], [1, 2]<] : ten>sor2x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
}< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
//
// %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가 있는 부동 소수점 요소 유형의 경우 op는 다음 IEEE-754 작업을 구현합니다.

• EQ: compareQuietEqual.
• NE: compareQuietNotEqual.
• GE: compareQuietGreaterEqual.
• GT: compareQuietGreater.
• LE: compareQuietLessEqual.
• LT: compareQuietLess.

compare_type = TOTALORDER가 있는 부동 소수점 요소 유형의 경우 작업은 IEEE-754의 totalOrder 및 compareQuietEqual 작업의 조합을 사용합니다.

복잡한 요소 유형의 경우 제공된 comparison_direction 및 compare_type를 사용하여 (real, imag) 쌍의 사전식 비교가 실행됩니다. 복소수에 순서를 적용하면 놀라운 의미가 발생하므로 향후 comparison_direction이 GE, GT, LE 또는 LT인 경우 복소수 지원을 삭제할 계획입니다(#560).

양자화된 유형의 경우 dequantize_compare(lhs, rhs, comparison_direction)를 실행합니다.

입력

출력

제약조건

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result)
• (C3) compare_type는 다음과 같이 정의됩니다.
  • is_signed_integer(element_type(lhs))인 경우 SIGNED
  • is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs))인 경우 UNSIGNED
  • is_float(element_type(lhs))인 경우 FLOAT 또는 TOTALORDER
  • is_complex(element_type(lhs))인 경우 FLOAT

예

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = <#stablehlocomparison_di>rection LT,
  compare_type = <#stablehlocomparison_>type FLOAT
}< : (t>ensor2xf<32, t>ens>or2xf32<) - >tensor2xi1
// %result: [true, false]

 예시 더보기

복잡한

시맨틱스

실수 값과 허수 값의 쌍인 lhs 및 rhs에서 복소수 값으로 요소별 변환을 실행하고 result 텐서를 생성합니다.

입력

출력

제약조건

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs)
• (C3) element_type(result)의 유형은 complex<E>이며 여기서 E = element_type(lhs)입니다.

예

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs)< : (t>ensor2xf<64, t>ens>or2xf64<) - tenso<r2x>>complexf64
// %result: [(1.0, 2.0), (3.0, 4.0)]

 예시 더보기

복합

시맨틱스

inputs 및 composite_attributes를 사용하고 results를 생성하는 다른 StableHLO 작업으로 구성된 작업을 캡슐화합니다. 작업의 의미는 decomposition 속성으로 구현됩니다. composite 작업은 프로그램 의미 체계를 변경하지 않고 분해로 대체할 수 있습니다. 분해를 인라인으로 처리해도 동일한 op 시맨틱이 제공되지 않는 경우 custom_call를 사용하는 것이 좋습니다.

version 필드 (기본값은 0)는 컴포지션의 의미가 변경되는 시점을 나타내는 데 사용됩니다.

입력

출력

제약조건

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

예

%results = "stablehlo.composite"(%input0, %input1) {
  name = "my_namespace.my_op",
  composite_attributes = {
    my_attribute = "my_value"
  },
  decomposition = @my_op,
 < ve>rsion = <1 :> i3>2
} : (<ten>sorf32, tensorf32) - tensorf32

 예시 더보기

이어붙이기

시맨틱스

지정된 인수와 동일한 순서로 dimension 차원을 따라 inputs을 연결하고 result 텐서를 생성합니다. 더 공식적으로는 result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1]입니다. 여기서 각 항목의 의미는 다음과 같습니다.

1. id = d0 + ... + dk-1 + kd.
2. d은 dimension과 같고 d0, ... 은 inputs의 d번째 차원 크기입니다.

입력

출력

제약조건

• (C1) same(element_type(inputs...)).
• (C2) dim(inputs..., dimension)를 제외한 same(shape(inputs...))
• (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
}< : (ten>sor3x2xi<64, ten>sor>1x2xi64<) - ten>sor4x2xi64
// %result: [[1, 2], [3, 4], [5, 6], [7, 8]]

 예시 더보기

상수

시맨틱스

상수 value에서 output 텐서를 생성합니다.

입력

출력

제약조건

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

예

%output = "stablehlo.constant"() {
  val<ue = dense[[0.0, 1.0], [>2.0, 3.0]<] : ten>sor2x2xf3>2
} : (<) - ten>sor2x2xf32
// %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로 변환됩니다. 복잡한 유형의 경우 아래를 참고하세요.

정수-정수, 정수-부동 소수점 또는 부동 소수점-부동 소수점 변환의 경우 소스 값을 대상 유형으로 정확하게 표현할 수 있으면 결과 값은 정확한 표현입니다. 그렇지 않으면 동작은 미정입니다(#180).

floating-point-to-integer 변환하는 경우 소수 부분이 잘립니다. 잘린 값을 대상 유형으로 나타낼 수 없는 경우 동작은 미정입니다 (#180).

복소수-복소수 변환은 실수부와 허수부를 변환하는 부동 소수점-부동 소수점 변환과 동일한 동작을 따릅니다.

complex-to-any-other-type 및 any-other-type-to-complex 변환의 경우 소스 허수 값이 무시되거나 대상 허수 값이 0으로 설정됩니다. 실수부의 변환은 부동 소수점 변환을 따릅니다.

원칙적으로 이 작업은 역양자화 (양자화된 텐서에서 일반 텐서로 변환), 양자화 (일반 텐서에서 양자화된 텐서로 변환), 재양자화 (양자화된 텐서 간 변환)를 표현할 수 있지만 현재는 첫 번째 사용 사례의 경우 uniform_dequantize, 두 번째 및 세 번째 사용 사례의 경우 uniform_quantize와 같은 전용 작업이 있습니다. 향후 이 두 작업은 convert로 병합될 수 있습니다 (#1576).

입력

출력

제약조건

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

예

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand)< : (t>ens>or3xi64<) - tenso<r3x>>complexf64
// %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]).
• j[d] = i[permutation[d]] 조건을 충족하는 permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1]

feature_group_count = 1 및 batch_group_count = 1인 경우 index_space(dim(result, output_spatial_dimensions...))의 모든 output_spatial_index에 대해 다음을 충족하는 result[result_shape(:, output_spatial_index, :)] = dot_product입니다.

• padding_value = constant(0, element_type(lhs)).
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1).
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides.
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations).
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true]). 이 기능은 사용되지 않는 것으로 보이므로 향후 삭제할 계획입니다 (#1181).
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension]).

feature_group_count > 1인 경우:

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

batch_group_count > 1인 경우:

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

양자화된 유형의 경우 dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result))를 실행합니다.

하이브리드 양자화 유형의 경우 hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs)를 실행합니다.

입력

출력

제약조건

• (C1) N = rank(lhs) = rank(rhs).
• (C2) size(window_strides) = N - 2
• (C3) 0 < window_strides.
• (C4) shape(padding) = [N - 2, 2]
• (C5) size(lhs_dilation) = N - 2.
• (C6) 0 < lhs_dilation
• (C7) size(rhs_dilation) = N - 2.
• (C8) 0 < rhs_dilation
• (C9) size(window_reversal) = N - 2.
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0.
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0.
• (C12) size(input_spatial_dimensions) = N - 2.
• (C13) input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]가 다음과 같이 주어집니다.
  • is_unique(input_dimensions).
  • 0 <= input_dimensions < N.
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0.
• (C17) size(kernel_spatial_dimensions) = N - 2.
• (C18) kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]:
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]가 제공된 경우:
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• (C21) 0 < feature_group_count.
• (C22) 0 < batch_group_count
• (C23) feature_group_count = 1 or batch_group_count = 1.
• (C24) size(precision_config) = 2.
• (C25) dim(result, result_dim)은 다음과 같이 정의됩니다.
  • result_dim = output_batch_dimension인 경우 dim(lhs, input_batch_dimension) / batch_group_count
  • result_dim = output_feature_dimension인 경우 dim(rhs, kernel_output_feature_dimension)
  • 그렇지 않으면 num_windows입니다.
  • output_spatial_dimensions[spatial_dim] = result_dim.
  • lhs_dim = input_spatial_dimensions[spatial_dim].
  • rhs_dim = kernel_spatial_dimensions[spatial_dim].
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1.
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1].
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1.
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim].
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1.
• (C26) rank(result) = N.
• 작업에서 양자화되지 않은 텐서를 사용하는 경우 다음을 충족해야 합니다.
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• 작업에서 양자화된 텐서를 사용하는 경우 다음을 충족해야 합니다.
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) is_per_axis_quantized(rhs)인 경우 quantization_dimension(rhs) = kernel_output_feature_dimension입니다.
  • (C30) is_per_axis_quantized(result)이면 quantization_dimension(result) = output_feature_dimension입니다.
  • is_quantized(lhs)인 경우:
  • (C31) storage_type(lhs) = storage_type(rhs)
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) is_per_tensor_quantized(rhs)이면 is_per_tensor_quantized(result)입니다.
  • !is_quantized(lhs)인 경우:
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result)

예

// %lhs: [[
//        [
//          [1], [2], [5], [6]
//        ],
//        [
//          [3], [4], [7], [8]
//        ],
//        [
//          [10], [11], [14], [15]
//        ],
//        [
//          [12], [13], [16], [17]
//        ]
//      ]]
//
// %rhs: [
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]]
//       ]
%result = "stablehlo.convolution"(%lhs, %rhs) {
  window_strid<es = arra>yi64: 4, 4,
  paddi<n>g = dense<0 : ten>sor2x2xi64,
  lhs_dilati<on = arra>yi64: 2, 2,
  rhs_dilati<on = arra>yi64: 1, 1,
  window_revers<al = arrayi1: fa>lse, false,
  // In the StableHLO dialect, dimension numbers are encoded vi<a:
  // `[input >dim<ensions]x[kernel >di>mensions]-[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" a<re spatial dimensions.
  d>imension_num>bers = #stablehlo.conv[b, 0, 1, f]x[0, 1, i, o]-[b, 0, 1, f],
  batch_group_count = 1 : i64,
  fea<ture_group_count >= 1 : i64,
 < precision_config> = [#stablehl<oprecision >DEFAULT,< #stablehlo>pre>cision <DEFAULT]
} >: (tensor1x4x4x1xi64, tensor3x3x1x1xi64) - tensor1x2x2x1xi64
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

 예시 더보기

코사인

시맨틱스

operand 텐서에 요소별 코사인 연산을 수행하여 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 부동 소수점의 경우: IEEE-754의 cos
• 복소수의 경우 복소수 코사인입니다.
• 양자화된 유형의 경우: dequantize_op_quantize(cosine, operand, type(result))

입력

출력

제약조건

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

예

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

 예시 더보기

count_leading_zeros

시맨틱스

operand 텐서에서 선행하는 0비트의 수를 요소별로 계산하고 result 텐서를 생성합니다.

입력

출력

제약조건

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

예

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

 예시 더보기

custom_call

시맨틱스

inputs 및 called_computations를 가져와 results를 생성하는 구현 정의 작업 call_target_name를 캡슐화합니다. has_side_effect, backend_config, api_version은 구현 정의 메타데이터를 추가로 제공하는 데 사용할 수 있습니다.

현재 이 작업에는 XLA 컴파일러에서 상응하는 작업의 유기적 진화를 반영하는 상당히 체계적이지 않은 메타데이터 모음이 포함되어 있습니다. 향후 이 메타데이터를 통합할 계획입니다(#741).

입력

출력

(XLA GPU 지원) 특수 custom_call 타겟

buffer 유형과 관련된 세 가지 특수 call_target_name가 있습니다. CreateBuffer는 초기화되지 않은 buffer를 만들고, Pin는 초기화된 buffer를 만들며, Unpin는 buffer를 할당 해제하고 buffer의 콘텐츠를 반환합니다.

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

%initialized_buffer = "stablehlo.custom_call"(%init_value) {
  call_target_name = "Pin&quo<t;,
 > ap>i_versi<on = >4 : i32,
} : (tensor4xf64) - memref4xf64

%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
  call_target_na<me = >&qu>ot;Unpi<n&quo>t;,
  api_version = 4 : i32,
} : (memref4xf64) - tensor4xf64

별칭

일부 custom_call 작업은 출력이 메모리를 공유하고 피연산자가 메모리를 공유해야 할 수 있습니다. 이는 output_operand_aliases를 통해 표현할 수 있습니다. 별칭 쌍 표현식은 출력 부분을 나타내는 출력 튜플 색인 목록과 피연산자 부분을 나타내는 피연산자 튜플 색인 목록과 함께 operand_index로 구성됩니다. 해당 유형이 tuple 유형이 아닌 경우 출력 또는 피연산자 튜플 인덱스 목록은 비어 있으며, 임의로 중첩된 튜플 유형의 경우 임의로 길 수 있습니다. 이는 XLA 별칭 표현과 유사합니다.

별칭 쌍의 출력 부분과 입력 부분은 유형이 동일해야 합니다. CreateBuffer, Pin, Unpin 호출이 아닌 custom_call 작업의 경우 buffer 피연산자는 최대 하나의 별칭 쌍에 표시될 수 있으며 buffer 출력은 하나의 별칭 쌍에 표시되어야 합니다.

예

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

%updated_buffer = "stablehlo.custom_call"(%buffer) {
  call_target_name = "Update",
  api_version = 4 : i32,
  output_operand_aliases< = [
    #stablehlo.output_operand_aliasoutput_tuple_indices = [],
      operand_ind>ex = 0,
     < oper>and>_tuple_<indic>es = []]
} : (memref4xf64) - memref4xf64

나누기

시맨틱스

피제수 lhs와 제수 rhs 텐서의 요소별 나눗셈을 실행하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

• 정수의 경우: 소수 부분이 삭제된 대수적 몫을 생성하는 정수 나눗셈입니다.
• 부동 소수점의 경우: IEEE-754의 division
• 복소수의 경우 복소수 나눗셈입니다.
• 양자화된 유형의 경우:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

입력

출력

제약조건

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

예

// %lhs: [17.1, -17.1, 17.1, -17.1]
// %rhs: [3.0, 3.0, -3.0, -3.0]
%result = "stablehlo.divide"(%lhs, %rhs)< : (t>ensor4xf<32, t>ens>or4xf32<) - t>ensor4xf32
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

 예시 더보기

dot_general

시맨틱스

lhs 슬라이스와 rhs 슬라이스 간의 내적을 계산하고 result 텐서를 생성합니다.

더 공식적으로는 result[result_index] = dot_product입니다. 여기서

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

양자화된 유형의 경우 dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result))를 실행합니다.

하이브리드 양자화 유형의 경우 hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs)를 실행합니다.

precision_config는 가속기 백엔드에서 계산의 속도와 정확도 간의 균형을 제어합니다. 다음 중 하나일 수 있습니다 (현재 이러한 enum 값의 의미는 명확하게 명시되어 있지 않지만 #755에서 이 문제를 해결할 예정임).

• DEFAULT: 가장 빠른 계산이지만 원래 숫자에 대한 근사치가 가장 정확하지 않습니다.
• HIGH: 계산이 느리지만 원래 숫자에 더 정확하게 근사합니다.
• HIGHEST: 가장 느린 계산이지만 원래 숫자에 가장 정확하게 근사합니다.

DotAlgorithm는 점 연산을 구현하는 데 사용되는 알고리즘의 기본 속성을 정의하며, 이는 정밀도도 정의합니다. 알고리즘 속성 필드가 설정된 경우 precision_config는 DEFAULT이어야 합니다. DotAlgorithms에는 기본값이 없습니다. 기본 매개변수는 구현에 따라 정의되기 때문입니다. 따라서 모든 점 알고리즘 필드를 None로 설정하여 빈 점 알고리즘을 지정할 수 있으며, 이 경우 precision_config 값이 대신 사용됩니다.

DotAlgorithm 필드에는 다음이 포함됩니다.

• lhs_precision_type 및 rhs_precision_type: 작업의 좌항과 우항이 반올림되는 정밀도입니다. 정밀도 유형은 입력 및 출력의 저장 유형과 독립적입니다.
• accumulation_type 누적에 사용되는 정밀도입니다.
• lhs_component_count, rhs_component_count, num_primitive_operations는 좌변 또는 우변을 여러 구성요소로 분해하고 이러한 값에 여러 '기본' 점 연산을 실행하는 알고리즘을 실행할 때 적용됩니다. 이는 일반적으로 더 높은 정밀도를 에뮬레이션하기 위한 것입니다 (예: 더 높은 정밀도 계산을 위해 bfloat16 인공지능 데이터 유형 활용: bf16_6x tf32_3x 등). 분해가 없는 알고리즘의 경우 이러한 값은 1으로 설정해야 합니다.
• allow_imprecise_accumulation: 일부 단계 (예: CUBLASLT_MATMUL_DESC_FAST_ACCUM)에서 낮은 정밀도로 누적이 허용되는지 지정합니다.

DotAlgorithm 속성의 예:

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


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


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

어떤 조합이 지원되는지는 구현에 달려 있습니다. 일반적으로 각 알고리즘이 StableHLO의 소비자에 의해 각 가속기 유형에서 지원된다고 보장되지는 않습니다. 지정된 알고리즘이 지원되지 않으면 대체 알고리즘으로 대체하는 대신 오류가 발생해야 합니다. StableHLO 검증은 최선을 다해 검증하여 모든 하드웨어에서 지원되지 않는 것으로 알려진 알고리즘을 방지합니다.

지원되는 알고리즘 값은 xla_data.proto > Algorithm을 참고하세요. 티켓 #2483에는 백엔드에서 지원하는 알고리즘에 관한 중앙 집중식 문서를 만들 계획이 나와 있습니다.

입력

출력

제약조건

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions).
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions)
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions).
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions)
• (C5) 0 <= lhs_batching_dimensions < rank(lhs).
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs)
• (C7) 0 <= rhs_batching_dimensions < rank(rhs).
• (C8) 0 <= rhs_contracting_dimensions < rank(rhs)
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...).
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...).
• (C11) size(precision_config) = 2.
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions).
• 작업에서 양자화되지 않은 텐서를 사용하는 경우 다음을 충족해야 합니다.
  • (C13) element_type(lhs) = element_type(rhs).
• 작업에서 양자화된 텐서를 사용하는 경우 다음을 충족해야 합니다.
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)
  • (C15) zero_points(rhs) = 0
  • (C16) is_per_axis_quantized(rhs)인 경우 quantization_dimension(rhs)이 rhs_contracting_dimensions에 없습니다.
  • is_quantized(lhs)인 경우:
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) is_per_tensor_quantized(rhs)이면 is_per_tensor_quantized(result)입니다.
  • !is_quantized(lhs)인 경우:
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation)인 경우:
  • (C21) precision_config... = DEFAULT.
  • (C22) 0 < lhs_component_count
  • (C23) 0 < rhs_component_count.
  • (C24) 0 < num_primitive_operations.

예

// %lhs: [
//        [[1, 2],
//         [3, 4]],
//        [[5, 6],
//         [7, 8]]
//       ]
// %rhs: [
//        [[1, 0],
//         [0, 1]],
//        [[1, 0],
//         [0, 1]]
//       ]
%result = "stablehlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #sta<blehlo.dot
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimension>s = [1]
  ,
  precision_config = [<#stablehloprecisi>on DEFAULT, <#stablehloprecisi>on DEFAULT],
  algorithm = #stablehlo.dot<_algorithm
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation >= false
  
}< : (tenso>r2x2x2xi<64, tenso>r2x>2x2xi64<) - tenso>r2x2x2xi64
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

 예시 더보기

dynamic_broadcast_in_dim

시맨틱스

이 작업은 broadcast_in_dim 작업과 기능적으로 동일하지만 결과 모양은 output_dimensions를 통해 동적으로 지정됩니다.

이 작업은 측정기준의 확장 동작에 관한 정적 지식을 표현하기 위해 선택적 속성 known_expanding_dimensions, known_nonexpanding_dimensions도 허용합니다. 지정하지 않으면 모든 측정기준이 확장될 수 있는 것으로 간주됩니다.

입력

출력

제약조건

• (C1) element_type(result)은 다음 식으로 주어집니다.
  • !is_per_axis_quantized(operand)인 경우 element_type(operand)
  • element_type(operand)(quantization_dimension(operand), scales(operand), zero_points(operand)이 각각 quantization_dimension(result), scales(result), zero_points(result)와 다를 수 있음)
• (C2) size(broadcast_dimensions) = rank(operand)
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions)
• (C5) axes(operand)의 모든 d에 대해 다음을 실행합니다.
  • dim(operand, d) = 1 또는
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • dim(operand, quantization_dimension(operand)) = 1이면 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))입니다.
• (C7) size(output_dimensions) = rank(result).
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions)
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand).

예

// %operand: [
//            [1, 2, 3]
//           ]
%operand = stablehlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = stablehlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "stablehlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
  broadcast_dimensio<ns = arra>yi64: 2, 1,
  known_expanding_dimensio<ns = a>rrayi64: 0,
  known_nonexpanding_dimensio<ns = a>rrayi64: 1
}< : (ten>sor1x3xi<64, t>ens>or3xi64<) - tenso>r2x3x2xi64
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 예시 더보기

dynamic_conv

시맨틱스

이 작업은 기능적으로 컨볼루션 작업과 동일하지만 패딩은 padding를 통해 동적으로 지정됩니다.

입력