StableHLO 사양

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 섹션은 프로그램 내에서 함께 실행되는 이러한 모든 연산의 의미 체계를 제공합니다. 마지막으로 표기법 섹션에서는 사양 전반에서 사용되는 표기법을 설명합니다.

이전 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) 모든 식별자에는 여러 종류의 식별자를 구분하는 시글(sigil)이 있고, 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입니다.

도형은 부분적으로 알 수 없거나 완전히 알 수 없는 (동적)일 수 있습니다. 예를 들어 tensor<?x2xf64>는 부분적으로 알 수 없고 tensor<?x?xf64>는 완전히 알 수 없습니다. 동적 크기 크기는 ?를 사용하여 표현됩니다. 도형의 순위는 지정 해제할 수 없습니다.

향후에는 크기 및 요소 유형 이상으로 텐서 유형을 확장하여 레이아웃(#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~C3), (C8)
storage_min 정수 상수 (C1), (C3), (C7)
storage_max 정수 상수 (C2), (C3), (C7)
expressed_type 부동 소수점 유형 (C4)
quantization_dimension 선택적 정수 상수 (C10-C12)
scales 부동 소수점 상수의 가변수 (C4-C6), (C9), (C10), (C13)
zero_points 정수 상수의 가변수 (C7-C9)

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

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

현재 QuantizationScale는 부동 소수점 상수이지만 승수 및 시프트로 표현되는 정수 기반 척도에 큰 관심이 있습니다. 조만간 이 기능을 살펴볼 계획입니다(#1404).

유형, 값, 양자화 텐서 유형에 1개의 제로 포인트만 있거나 잠재적으로 여러 개의 제로 포인트가 있을 수 있는지 여부 등 QuantizationZeroPoint의 의미 체계에 관한 논의가 진행 중입니다. 이 설명의 결과에 따라 0포인트에 관한 사양은 향후 변경될 수 있습니다. (#1405)

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

마지막으로, 알 수 없는 크기 크기를 나타낼 계획 (#1407)과 마찬가지로 알 수 없는 배율과 0점을 나타내는 방법을 알아볼 계획입니다.

양자화 텐서 유형은 양자화된 요소가 있는 텐서를 나타냅니다. 이러한 텐서는 요소에 일반 요소 유형 대신 양자화된 요소 유형이 있다는 점을 제외하고 일반 텐서와 정확히 동일합니다.

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

  • 텐서당 양자화:
    • 추가 제약 조건이 없습니다.
  • 축당 양자화:
    • (C13) quantization_dimension < rank(self).
    • (C14) dim(self, quantization_dimension) = size(scales)입니다.
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차원 텐서 값으로 표현하는 것이 관용적입니다.

  • 부울 유형은 불리언 값 truefalse를 나타냅니다.
  • 정수 유형은 부호 있는 유형 (si) 또는 부호 없는 유형 (ui)일 수 있으며 지원되는 비트 너비 (4, 8, 16, 32 또는 64) 중 하나를 갖습니다. 부호 있는 siN 유형은 -2^(N-1) 이상 2^(N-1)-1까지의 정수 값을 나타내고, 부호 없는 uiN 유형은 0부터 2^N-1까지의 정수 값을 나타냅니다.
  • 부동 소수점 유형은 다음 중 하나일 수 있습니다.
  • 복합 유형은 동일한 요소 유형실수허수부를 갖는 복합 값을 나타냅니다. 지원되는 복합 유형은 complex<f32> (두 부분 모두 f32 유형) 및 complex<f64>(두 부분 모두 f64 유형)입니다.
FunctionType ::= '(' InputTypes ')' '->' '(' OutputTypes ')'
InputTypes ::= [ValueType {',' ValueType}]
OutputTypes ::= [ValueType {',' ValueType}]

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

StringType ::= 'string'

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

운영

안정적 HLO 작업 (작업이라고도 함)은 머신러닝 모델에서 상위 수준 작업들의 폐쇄형 집합을 나타냅니다. 위에서 설명한 것처럼 StableHLO 구문은 MLIR에서 큰 영향을 받으며, 이는 꼭 가장 인체공학적인 대안은 아니지만 ML 프레임워크와 ML 컴파일러 간의 상호 운용성을 개선하려는 StableHLO의 목표에 가장 적합합니다.

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

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

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

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

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

입력 함수 (익명 함수라고도 함)는 이름이 지정된 함수와 매우 유사합니다. 단, 1) 식별자가 없으므로 (따라서 이름이 '익명'임) 2) 출력 유형을 선언하지 않습니다 (출력 유형은 함수 내의 return 연산에서 추론됨).

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

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

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

현재 실제 상황에서의 StableHLO 프로그램에는 이 문서에서 설명하지 않는 속성이 포함되어 있는 경우가 있습니다. 향후 Google은 이러한 속성을 StableHLO opset에 흡수하거나 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'

부울 상수는 불리언 값 truefalse를 나타냅니다. 불리언 상수의 유형은 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]

예시 더보기

추가

시맨틱

두 텐서 lhsrhs를 요소별로 추가하고 result 텐서를 생성합니다. 요소 유형에 따라 다음을 실행합니다.

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

입력

라벨 이름 유형 제약조건
(I1) lhs 텐서 또는 양자화 텐서 (C1~C6)
(I2) rhs 텐서 또는 양자화 텐서 (C1~C5), (C7)

출력

이름 유형 제약조건
result 텐서 또는 양자화 텐서 (C1~C7)

제약조건

  • 연산에서 양자화되지 않은 텐서를 사용하는 경우:
    • (C1) type(lhs) = type(rhs) = type(result).
  • 연산에서 양자화 텐서를 사용하는 경우:
    • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result).
    • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result)
    • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)입니다.
    • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result)입니다.
    • (C6) is_per_axis_quantized(lhs)인 경우 quantization_dimension(lhs) = quantization_dimension(result)
    • (C7) is_per_axis_quantized(rhs)인 경우 quantization_dimension(rhs) = quantization_dimension(result)

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

예시 더보기

after_all

시맨틱

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

입력

라벨 이름 유형
(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 프로세스 그리드의 각 프로세스 그룹 내에서 각 프로세스의 operand 텐서 값을 all_gather_dim에 연결하여 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(tensor<E>, tensor<E>) -> (tensor<E>) 유형(여기서 is_promotable(element_type(operand), E))입니다.
  • (C6) shape(result) = shape(operand)
  • (C7) element_type(result) = E

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [1, 2, 3, 4]
// %operand@(1, 0): [5, 6, 7, 8]
%result = "stablehlo.all_reduce"(%operand) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<4xi64>) -> tensor<4xi64>
// %result@(0, 0): [6, 8, 10, 12]
// %result@(1, 0): [6, 8, 10, 12]

예시 더보기

all_to_all

시맨틱

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]]

예시 더보기

시맨틱

두 텐서 lhsrhs에 요소별 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

시맨틱

lhsrhs 텐서에서 요소별 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_scalegrad_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, varianceresultbaseline_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_varoutputbaseline_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, :, :]는 하부삼각형(lowertrue인 경우) 또는 상삼각형 (lowerfalse인 경우) 중 하나의 형태로 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(is_quantized(result) ? quantize(0, element_type(result)) : 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_directioncompare_type에 따라 lhsrhs 텐서의 요소별 비교를 수행하고 result 텐서를 생성합니다.

comparison_directioncompare_type의 값은 다음과 같은 시맨틱스를 갖습니다.

불리언 및 정수 요소 유형:

  • EQ: lhs = rhs.
  • NE: lhs != rhs.
  • GE: lhs >= rhs.
  • GT: lhs > rhs.
  • LE: lhs <= rhs.
  • LT: lhs < rhs.

compare_type = FLOAT가 있는 부동 소수점 요소 유형의 경우 작업은 다음과 같은 IEEE-754 연산을 구현합니다.

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

compare_type = TOTALORDER이 있는 부동 소수점 요소 유형의 경우 연산은 IEEE-754의 totalOrder 연산과 compareQuietEqual 연산을 함께 사용합니다.

복잡한 요소 유형의 경우 제공된 comparison_directioncompare_type를 사용하여 (real, imag) 쌍의 사전식 비교가 실행됩니다. 복소수에 정렬을 적용하려면 예기치 않은 의미 체계가 포함되므로 향후 comparison_directionGE, 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입니다.
    • is_float(element_type(lhs))인 경우 FLOAT 또는 TOTALORDER입니다.
    • is_complex(element_type(lhs))인 경우 FLOAT입니다.

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

예시 더보기

복잡함

시맨틱

실수값과 허수값 쌍(lhsrhs)에서 복소수 값으로 요소별 변환을 수행하고 result 텐서를 생성합니다.

입력

라벨 이름 유형 제약조건
(I1) lhs f32 또는 f64 유형의 텐서 (C1~C3)
(I2) rhs f32 또는 f64 유형의 텐서 (C1)

출력

이름 유형 제약조건
result 복합 유형의 텐서 (C2), (C3)

제약조건

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

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

예시 더보기

복합

시맨틱

다른 StableHLO 작업으로 구성된 (구성된) 작업을 캡슐화하여 inputscomposite_attributes를 가져와 results를 생성합니다. 작업의 시맨틱은 decomposition 속성으로 구현됩니다. composite 작업은 프로그램 시맨틱을 변경하지 않고 분해로 대체할 수 있습니다. 인라인 처리에서 동일한 연산 시맨틱스를 제공하지 않는 경우 custom_call를 사용하는 것이 좋습니다.

version 필드 (기본값 0)는 복합의 의미 체계가 변경되는 시점을 나타내는 데 사용됩니다.

입력

라벨 이름 유형
(I1) inputs 가변 값 수
(I2) name string 유형의 상수
(I3) composite_attributes 속성 사전
(I4) decomposition string 유형의 상수
(I5) version si32 유형의 상수

출력

이름 유형
results 가변 값 수

제약조건

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

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

예시 더보기

concatenate

시맨틱

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

  1. id = d0 + ... + dk-1 + kd.
  2. ddimension와 같고 d0inputsd번째 차원 크기입니다.

입력

라벨 이름 유형 제약조건
(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-typeany-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를 생성합니다. 다음 다이어그램은 구체적인 예를 사용하여 result의 요소가 lhsrhs에서 계산되는 방식을 보여줍니다.

컨볼루션

더 공식적으로, lhs의 기간을 표현할 수 있도록 lhs의 관점에서 다음과 같이 입력의 재구성을 고려하세요.

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

이 재구성에는 다음과 같은 도우미 함수가 사용됩니다.

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

feature_group_count = 1이고 batch_group_count = 1인 경우 index_space(dim(result, output_spatial_dimensions...))의 모든 output_spatial_index에 대해 result[result_shape(:, output_spatial_index, :)] = dot_product이며 각 항목의 의미는 다음과 같습니다.

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

feature_group_count > 1인 경우:

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

batch_group_count > 1인 경우:

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

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

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

입력

라벨 이름 유형 제약조건
(I1) lhs 텐서 또는 텐서당 양자화 텐서 (C1), (C10-C11), (C14) (C25), (C27-C28), (C31-C32), (C34)
(I2) rhs 텐서 또는 양자화 텐서 (C1), (C14-C16), (C25), (C27-C29), (C31-C34)
(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), (C29)
(I13) kernel_spatial_dimensions si64 유형의 1차원 텐서 상수 (C17-C18), (C25)
(I14) output_batch_dimension si64 유형의 상수 (C20), (C25)
(I15) output_feature_dimension si64 유형의 상수 (C20), (C25), (C30)
(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), (C32-34)

제약조건

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

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

예시 더보기

코사인

시맨틱

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

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

입력

라벨 이름 유형 제약조건
(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

시맨틱

inputscalled_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))를 실행합니다.

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

precision_config는 가속기 백엔드에서 계산의 속도와 정확성 간의 균형을 제어합니다. 다음 중 하나일 수 있습니다. 현재는 이러한 열거형 값의 의미 체계가 충분히 지정되지 않지만 #755에서 이 문제를 해결할 계획입니다.

  • DEFAULT: 가장 빠른 계산이지만 원래 숫자에 대한 근사값은 가장 낮습니다.
  • HIGH: 계산은 느리지만 원래 숫자에 더 가깝습니다.
  • HIGHEST: 계산은 가장 느리지만 원래 숫자에 가장 정확합니다.

입력

라벨 이름 유형 제약조건
(I1) lhs 텐서 또는 텐서당 양자화 텐서 (C5-C6), (C9-C10), (C12-C14), (C17-C18), (C20)
(I2) rhs 텐서 또는 양자화 텐서 (C7-C10), (C12-C20)
(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), (C16)
(I7) precision_config DEFAULT, HIGH, HIGHEST의 다양한 enum 수 (C11)

출력

이름 유형 제약조건
result 텐서 또는 양자화 텐서 (C12), (C14), (C18-C20)

제약조건

  • (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)

// %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_broadcast_in_dim

시맨틱

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

이 작업은 선택적 속성인 known_expanding_dimensions, known_non_expanding_dimensions도 허용하여 크기의 확장 동작에 관한 정적 지식을 표현합니다. 지정하지 않으면 모든 크기가 확장 가능한 것으로 간주됩니다.

입력

라벨 이름 유형 제약조건
(I1) operand 텐서 또는 양자화 텐서 (C1-C2), (C5-C6), (C9)
(I2) output_dimensions 정수 유형의 1차원 텐서 (C7)
(I3) broadcast_dimensions 정수 유형의 1차원 상수 텐서 (C2-C6)
(I4) known_expanding_dimensions 정수 유형의 1차원 상수 텐서 (C8-C9)
(I5) known_non_expanding_dimensions 정수 유형의 1차원 상수 텐서 (C8-C9)

출력

이름 유형 제약조건
result 텐서 또는 양자화 텐서 (C1), (C3), (C5-C7)

제약조건

  • (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)))입니다.
  • (C7) size(output_dimensions) = rank(result)
  • (C8) is_unique(known_expanding_dimensions + known_non_expanding_dimensions)
  • (C9) 0 <= known_expanding_dimensions < rank(operand)입니다.
  • (C10) 0 <= known_non_expanding_dimensions < rank(operand)

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

예시 더보기

dynamic_conv

시맨틱

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

입력

라벨 이름 유형 제약조건
(I1) lhs 텐서 또는 텐서당 양자화 텐서 (C1), (C10-C11), (C14) (C25), (C26-C27), (C30-C31), (C33)
(I2) rhs 텐서 또는 양자화 텐서 (C1), (C14-C16), (C26-C28), (C30-C33)
(I3) padding 정수 유형의 2차원 텐서 (C4)
(I4) window_strides si64 유형의 1차원 텐서 상수 (C2-C3)
(I5) lhs_dilation si64 유형의 1차원 텐서 상수 (C5-C6)
(I6) rhs_dilation si64 유형의 1차원 텐서 상수 (C7-C8)
(I7) window_reversal i1 유형의 1차원 텐서 상수 (C9)
(I8) input_batch_dimension si64 유형의 상수 (C10), (C13)
(I9) input_feature_dimension si64 유형의 상수 (C11), (C13-C14)
(I10) input_spatial_dimensions si64 유형의 1차원 텐서 상수 (C12), (C13)
(I11) kernel_input_feature_dimension si64 유형의 상수 (C14), (C18)
(I12) kernel_output_feature_dimension si64 유형의 상수 (C15-C16), (C18), (C28)
(I13) kernel_spatial_dimensions si64 유형의 1차원 텐서 상수 (C17-C18)
(I14) output_batch_dimension si64 유형의 상수 (C20)
(I15) output_feature_dimension si64 유형의 상수 (C20), (C29)
(I16) output_spatial_dimensions si64 유형의 1차원 텐서 상수 (C19-C20)
(I17) feature_group_count si64 유형의 상수 (C11), (C14), (C16), (C21), (C23)
(I18) batch_group_count si64 유형의 상수 (C10), (C15), (C22), (C23)
(I19) precision_config DEFAULT, HIGH, HIGHEST의 다양한 enum 수 (C24)

출력

이름 유형 제약조건
result 텐서 또는 양자화 텐서 (C25-C27), (C29), (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)입니다.
  • 연산에서 양자화 텐서를 사용하는 경우:
    • (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]]]
//        ]
// %padding: [[1, 1],
//            [1, 1]]
%result = "stablehlo.dynamic_conv"(%lhs, %rhs, %padding) {
  window_strides = array<i64: 4, 4>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  dimension_numbers = #stablehlo.conv<raw
    input_batch_dimension = 0,
    input_feature_dimension = 3,
    input_spatial_dimensions = [0, 1],
    kernel_input_feature_dimension = 2,
    kernel_output_feature_dimension = 3,
    kernel_spatial_dimensions = [0, 1],
    output_batch_dimension = 0,
    output_feature_dimension = 3,
    output_spatial_dimensions = [1, 2]
  >,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>, tensor<2x2xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[1], [5]],
//            [[10], [14]]
//          ]]

예시 더보기

dynamic_gather

시맨틱

이 작업은 gather 작업과 기능적으로 동일하며 slice_sizes가 동적으로 값으로 지정됩니다.

입력

라벨 이름 유형 제약조건
(I1) operand 텐서 또는 텐서당 양자화 텐서 (C1), (C7), (C10-C12), (C14)
(I2) start_indices 정수 유형의 텐서 (C2), (C3), (C13)
(I3) slice_sizes 정수 유형의 1차원 텐서 (C8), (C11-C13)
(I4) offset_dims si64 유형의 1차원 텐서 상수 (C1), (C4-C5), (C13)
(I5) collapsed_slice_dims si64 유형의 1차원 텐서 상수 (C1), (C6-C8), (C13)
(I6) start_index_map si64 유형의 1차원 텐서 상수 (C3), (C9), (C10)
(I7) index_vector_dim si64 유형의 상수 (C2), (C3), (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의 크기 크기가 포함되지 않는다는 점이 다릅니다.
    • combinebatch_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]]
//                 ]
// %slize_sizes: [1, 2, 2]
%result = "stablehlo.dynamic_gather"(%operand, %start_indices, %slize_sizes) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [1, 0],
    index_vector_dim = 2>,
  indices_are_sorted = false
} : (tensor<3x4x2xi64>, tensor<2x3x2xi64>, tensor<3xi64>) -> tensor<2x3x2x2xi64>
// %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]]
//            ]
//          ]

예시 더보기

dynamic_iota

시맨틱

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

입력

라벨 이름 유형 제약조건
(I1) output_shape 정수 유형의 1차원 텐서 (C1), (C2)
(I2) iota_dimension si64 (C1)

출력

이름 유형 제약조건
result 정수, 부동 소수점, 복합 유형 또는 텐서당 양자화 텐서의 텐서 (C2)

제약조건

  • (C1) 0 <= iota_dimension < size(output_shape).
  • (C2) rank(result) = size(output_shape).

%output_shape = stablehlo.constant dense<[4, 5]> : tensor<2xi64>
%result = "stablehlo.dynamic_iota"(%output_shape) {
  iota_dimension = 0 : i64
} : (tensor<2xi64>) -> tensor<4x5xi64>
// %result: [
//           [0, 0, 0, 0, 0],
//           [1, 1, 1, 1, 1],
//           [2, 2, 2, 2, 2],
//           [3, 3, 3, 3, 3]
//          ]

예시 더보기

dynamic_pad

시맨틱

이 작업은 패드 작업과 기능적으로 동일하지만 edge_padding_low, edge_padding_high, interior_padding가 동적으로 값으로 지정됩니다.

입력

라벨 이름 유형 제약조건
(I1) operand 텐서 또는 텐서당 양자화 텐서 (C1), (C2), (C4)
(I2) padding_value 0차원 텐서 또는 텐서당 양자화 텐서 (C1)
(I3) edge_padding_low 정수 유형의 1차원 텐서 (C1), (C4)
(I4) edge_padding_high 정수 유형의 1차원 텐서 (C1), (C4)
(I5) interior_padding 정수 유형의 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
// %edge_padding_low: [0, 1]
// %edge_padding_high: [2, 1]
// %interior_padding: [1, 2]
%result = "stablehlo.dynamic_pad"(%operand, %padding_value,
  %edge_padding_low, %edge_padding_high, %interior_padding
) : (tensor<2x3xi64>, tensor<i64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<5x9xi64>
// %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]
//          ]

예시 더보기

dynamic_reshape

시맨틱

이 작업은 reshape 연산과 기능적으로 동일하지만 결과 형태는 output_shape를 통해 동적으로 지정됩니다.

입력

라벨 이름 유형 제약조건
(I1) operand 텐서 또는 양자화 텐서 (C1~C3)
(I2) output_shape 정수 유형의 1차원 텐서 (C4)

출력

이름 유형 제약조건
result 텐서 또는 양자화 텐서 (C1~C4)

제약조건

  • (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).
  • (C4) size(output_shape) = rank(result)입니다.

// %operand: [[1, 2, 3], [4, 5, 6]]
// %output_shape: [3, 2]
%result = "stablehlo.dynamic_reshape"(%operand, %output_shape) : (tensor<2x3xi64>, tensor<2xi64>) -> tensor<3x2xi64>
// %result: [[1, 2], [3, 4], [5, 6]]

예시 더보기

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 = array<i64: 2, 2>
} : (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 텐서에 요소별 지수 빼기 연산을 수행하고 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의 경우 resultL = 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의 경우 resultfft_type = FFT 계산의 역으로 정의됩니다. 예를 들어 L = 3의 경우 다음과 같습니다.

  • result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]).
  • result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :] = ifft(result2[i0, ..., :]).

또한 rfft 함수가 부동 소수점 유형의 1차원 텐서를 취하는 경우 동일한 부동 소수점 의미 체계의 복합 유형의 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의 경우 resultL = 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의 경우 resultfft_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) operandresult 요소 유형 간의 관계가 다음과 같이 다릅니다.
    • 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) operandresult 사이에 부동 소수점 유형의 텐서 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 = array<i64: 4>
} : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
// %result: [(1.0, 0.0), (1.0, 0.0), (1.0, 0.0), (1.0, 0.0)]

바닥

시맨틱

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]: 여기서 bibatch_index의 개별 요소이고, index_vector_dimrank(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입니다.
  • axes(operand)d_operand,
    • d_operand = operand_batching_dims[i_batching]d_start = start_indices_batching_dims[i_batching]인 경우 full_batching_index[d_operand] = batch_index[d_start - (d_start < index_vector_dim ? 0 : 1)]입니다.
    • 그 밖의 경우에는 full_batching_index[d_operand] = 0입니다.
  • offset_index = result_index[offset_dims...].
  • full_offset_index = [oi0, ..., 0, ..., oiN]. 여기서 oioffset_index의 개별 요소이고 0collapsed_slice_dimsoperand_batching_dims의 색인에 삽입됩니다.
  • operand_index = full_start_index + full_batching_index + full_offset_index.

indices_are_sortedtrue이면 구현에서는 start_indicesstart_index_map을 기준으로 정렬되었다고 가정할 수 있습니다. 그렇지 않으면 동작이 정의되지 않습니다. 더 공식적으로는 indices(result)의 모든 i1 < i2에 대해 full_start_index(i1) <= full_start_index(i2)입니다.

입력

라벨 이름 유형 제약조건
(I1) operand 텐서 또는 텐서당 양자화 텐서 (C1), (C8), (C11), (C17), (C19-C21), (C23)
(I2) start_indices 정수 유형의 텐서 (C2-C3), (C14), (C17), (C22)
(I3) offset_dims si64 유형의 1차원 텐서 상수 (C1), (C4-C5), (C22)
(I4) collapsed_slice_dims si64 유형의 1차원 텐서 상수 (C1), (C6-C9), (C22)
(I5) operand_batching_dims si64 유형의 1차원 텐서 상수 (C1), (C6), (C10-C12), (C16-C18), (C22)
(I6) start_indices_batching_dims si64 유형의 1차원 텐서 상수 (C13-C17)
(I7) start_index_map si64 유형의 1차원 텐서 상수 (C3), (C18-C19)
(I8) index_vector_dim si64 유형의 상수 (C2-C3), (C15), (C22)
(I9) slice_sizes si64 유형의 1차원 텐서 상수 (C9), (C12), (C20-C22)
(I10) indices_are_sorted i1 유형의 상수

출력

이름 유형 제약조건
result 텐서 또는 텐서당 양자화 텐서 (C5), (C22-C23)

제약조건

  • (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims) + size(operand_batching_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(concatenate(collapsed_slice_dims, operand_batching_dims))
  • (C7) is_sorted(collapsed_slice_dims)
  • (C8) 0 <= collapsed_slice_dims < rank(operand)
  • (C9) slice_sizes[collapsed_slice_dims...] <= 1입니다.
  • (C10) is_sorted(operand_batching_dims)
  • (C11) 0 <= operand_batching_dims < rank(operand).
  • (C12) slice_sizes[operand_batching_dims...] <= 1.
  • (C13) is_unique(start_indices_batching_dims).
  • (C14) 0 <= start_indices_batching_dims < rank(start_indices)입니다.
  • (C15) index_vector_dim not in start_indices_batching_dims입니다.
  • (C16) size(operand_batching_dims) == size(start_indices_batching_dims)입니다.
  • (C17) dim(operand, operand_batching_dims...) = dim(start_indices, start_indices_batching_dims...)입니다.
  • (C18) is_unique(concatenate(start_index_map, operand_batching_dims))입니다.
  • (C19) 0 <= start_index_map < rank(operand)입니다.
  • (C20) size(slice_sizes) = rank(operand)
  • (C21) 0 <= slice_sizes <= shape(operand).
  • (C22) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) 각 항목의 의미는 다음과 같습니다.
    • batch_dim_sizes = shape(start_indices). 단, index_vector_dim에 해당하는 start_indices의 크기 크기는 포함되지 않습니다.
    • offset_dim_sizes = slice_sizes. 단, collapsed_slice_dimsoperand_batching_dims에 해당하는 slice_sizes의 크기 크기는 포함되지 않습니다.
    • combinebatch_dims에 해당하는 축에 batch_dim_sizes를 배치하고 offset_dims에 해당하는 축에 offset_dim_sizes를 배치합니다.
  • (C23) 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]]
//            ],
//            [
//             [[25, 26], [27, 28], [29, 30], [31, 32]],
//             [[33, 34], [35, 36], [37, 38], [39, 40]],
//             [[41, 42], [43, 44], [45, 46], [47, 48]]
//            ]
//           ]
// %start_indices: [
//                  [
//                   [[0, 0], [1, 0], [2, 1]],
//                   [[0, 1], [1, 1], [0, 9]]
//                  ],
//                  [
//                   [[0, 0], [2, 1], [2, 2]],
//                   [[1, 2], [0, 1], [1, 0]]
//                  ]
//                 ]
%result = "stablehlo.gather"(%operand, %start_indices) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [3, 4],
    collapsed_slice_dims = [1],
    operand_batching_dims = [0],
    start_indices_batching_dims = [1],
    start_index_map = [2, 1],
    index_vector_dim = 3>,
  slice_sizes = array<i64: 1, 1, 2, 2>,
  indices_are_sorted = false
} : (tensor<2x3x4x2xi32>, tensor<2x2x3x2xi64>) -> tensor<2x2x3x2x2xi32>
// %result: [
//           [
//            [
//             [[1, 2], [3, 4]],
//             [[3, 4], [5, 6]],
//             [[13, 14], [15, 16]]
//            ],
//            [
//             [[33, 34], [35, 36]],
//             [[35, 36], [37, 38]],
//             [[41, 42], [43, 44]]
//            ]
//           ],
//           [
//            [
//             [[1, 2], [3, 4]],
//             [[13, 14], [15, 16]],
//             [[21, 22], [23, 24]]
//            ],
//            [
//             [[43, 44], [45, 46]],
//             [[33, 34], [35, 36]],
//             [[27, 28], [29, 30]]
//            ]
//           ]
//          ]

예시 더보기

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))
  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는 먼저 나오는 페이로드 값과 마지막에 오는 토큰으로 구성됩니다. 앞으로는 명확성을 개선하기 위해 페이로드와 토큰을 개별 출력 2개로 분할할 계획입니다. (#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[output_index] = constant(is_quantized(output) ? quantize(output_index[iota_dimension], element_type(output)) : output_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 텐서에서 요소별 로그 + 연산 1개를 실행하고 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와 함께 매핑 함수 computationinputs에 적용하고 result 텐서를 생성합니다.

더 공식적으로는 result[result_index] = computation(inputs...[result_index])입니다.

입력

라벨 이름 유형 제약조건
(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 = array<i64: 0, 1>
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[0, 5], [12, 21]]

예시 더보기

최대

시맨틱

텐서 lhsrhs에서 요소별 max 연산을 실행하고 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]]

예시 더보기

최소

시맨틱

텐서 lhsrhs에 요소별 최소 연산을 실행하고 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]]

예시 더보기

곱하기

시맨틱

두 텐서 lhsrhs의 요소별 곱을 실행하고 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

예시 더보기

사용하여

시맨틱

두 텐서 lhsrhs에 대해 요소별 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"(%input0, %token) {
  outfeed_config = ""
} : (tensor<2x2x2xi64>, !stablehlo.token) -> !stablehlo.token

예시 더보기

패드

시맨틱

지정된 padding_value가 있는 텐서 주위와 텐서 요소 사이의 패딩으로 operand를 확장합니다.

edge_padding_lowedge_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 = array<i64: 0, 1>,
  edge_padding_high = array<i64: 2, 1>,
  interior_padding = array<i64: 1, 2>
} : (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_transfertrue이면 작업은 호스트에서 데이터를 전송합니다. 그렇지 않으면 다른 기기에서 데이터를 전송합니다. 이는 구현에 의해 정의됩니다. 이 플래그는 channel_type에 제공된 정보를 복제하므로 향후에는 이 중 하나만 유지할 계획입니다. (#666)

results는 먼저 나오는 페이로드 값과 마지막에 오는 토큰으로 구성됩니다. 앞으로는 명확성을 개선하기 위해 페이로드와 토큰을 개별 출력 2개로 분할할 계획입니다. (#670)

입력

라벨 이름 유형 제약조건
(I1) token token (C4)
(I2) channel_id si64 유형의 상수
(I3) channel_type DEVICE_TO_DEVICEHOST_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

시맨틱

축소 함수 bodydimensions와 함께 inputsinit_values에 적용하고 results 텐서를 생성합니다.

축소 순서는 구현을 통해 정의됩니다. 즉, bodyinit_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)에 있는 모든 indexinput_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 = array<i64: 1>
} : (tensor<1x6xi64>, tensor<i64>) -> tensor<1xi64>
// %result = [15]

예시 더보기

reduce_precision

시맨틱

exponent_bitsmantissa_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

시맨틱

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) computationis_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

시맨틱

축소 함수 bodyinputsinit_values의 윈도우에 적용하고 results를 생성합니다.

다음 다이어그램은 구체적인 예를 사용하여 inputs...에서 results...의 요소를 계산하는 방법을 보여줍니다.

reduce_window

보다 공식적으로 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(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) 유형(is_promotable(element_type(inputs[i]), Ei))입니다.
  • (C14) same(shape(results...))입니다.
  • (C15) shape(results[0]) = num_windows 각 항목의 의미는 다음과 같습니다.
    • dilated_input_shape = shape(inputs[0]) = 0 ? 0 : (shape(inputs[0]) - 1) * base_dilations + 1.
    • padded_input_shape = padding[:, 0] + dilated_input_shape + padding[:, 1].
    • dilated_window_shape = (window_dimensions - 1) * window_dilations + 1.
    • is_empty_window = padded_input_shape = 0 || dilated_window_shape > padded_input_shape.
    • num_windows = is_empty_window ? 0 : floor((padded_input_shape - dilated_window_shape) / window_strides) + 1.
  • (C16) [0,N)의 모든 i에 대해 element_type(results[i]) = Ei.

// %input = [[1, 2], [3, 4], [5, 6]]
// %init_value = 0
%result = "stablehlo.reduce_window"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  window_dimensions = array<i64: 2, 1>,
  window_strides = array<i64: 4, 1>,
  base_dilations = array<i64: 2, 1>,
  window_dilations = array<i64: 3, 1>,
  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_indexoperand_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 = array<i64: 1>
} : (tensor<3x2xi32>) -> tensor<3x2xi32>
// %result: [[2, 1], [4, 3], [6, 5]]

예시 더보기

RNG

시맨틱

rng_distribution 알고리즘을 사용하여 랜덤 숫자를 생성하고 지정된 형태 shaperesult 텐서를 생성합니다.

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 UNIFORMNORMAL의 열거형 (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입니다.
    • rng_algorithm = PHILOX인 경우 2 또는 3입니다.

// %initial_state: [1, 2]
%output_state, %output = "stablehlo.rng_bit_generator"(%initial_state) {
  rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
} : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x2xui64>)
// %output_state: [1, 6]
// %output: [
//           [9236835810183407956, 16087790271692313299],
//           [18212823393184779219, 2658481902456610144]
//          ]

round_nearest_afz

시맨틱

operand 텐서에서 가장 가까운 정수를 향해 요소별 반올림을 수행하고 0으로부터의 연결을 해제하고 result 텐서를 생성합니다. IEEE-754 사양에서 roundToIntegralTiesToAway 작업을 구현합니다. 양자화 유형의 경우 dequantize_op_quantize(round_nearest_afz, operand, type(result))를 실행합니다.

입력

라벨 이름 유형 제약조건
(I1) operand 부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서 (C1)

출력

이름 유형 제약조건
result 부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서 (C1)

제약조건

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

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_afz"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-3.0, 0.0, 1.0, 1.0, 3.0]

예시 더보기

round_nearest_even

시맨틱

operand 텐서에서 가장 가까운 정수를 향해 요소별 반올림을 수행하고 짝수 정수와의 연결을 끊고 result 텐서를 생성합니다. IEEE-754 사양에서 roundToIntegralTiesToEven 작업을 구현합니다. 양자화 유형의 경우 dequantize_op_quantize(round_nearest_even, operand, type(result))를 실행합니다.

입력

라벨 이름 유형 제약조건
(I1) operand 부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서 (C1)

출력

이름 유형 제약조건
result 부동 소수점 유형 또는 텐서당 양자화 텐서의 텐서 (C1)

제약조건

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

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_even"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-2.0, 0.0, 0.0, 1.0, 2.0]

예시 더보기

RQRT

시맨틱

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... 색인을 자세히 설명합니다.

scatter

더 공식적으로 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], 여기서 siupdate_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입니다.
  • axes(inputs[0])d_input,
    • d_input = input_batching_dims[i_batching]d_start = scatter_indices_batching_dims[i_batching]인 경우 full_batching_index[d_input] = update_scatter_index[d_start - (d_start < index_vector_dim ? 0 : 1)]입니다.
    • 그 밖의 경우에는 full_batching_index[d_input] = 0입니다.
  • update_window_index = update_index[update_window_dims...].
  • full_window_index = [wi0, ..., 0, ..., wiN]. 여기서 wiupdate_window_index의 개별 요소이고 0inserted_window_dimsinput_batching_dims의 색인에 삽입됩니다.
  • result_index = full_start_index + full_batching_index + full_window_index.

이 경우 results = exec(schedule, inputs)이며 각 항목의 의미는 다음과 같습니다.

  • scheduleindex_space(updates[0])의 구현 정의 순열입니다.
  • exec([update_index, ...], results) = exec([...], updated_results) 각 항목의 의미는 다음과 같습니다.
    • result_indexshape(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_resultsresults...[result_index]updated_values...로 설정된 results의 사본입니다.
    • 그렇지 않은 경우 다음 단계를 따릅니다.
    • updated_results = results.
  • exec([], results) = results.

indices_are_sortedtrue이면 구현에서는 scatter_indicesscatter_dims_to_operand_dims을 기준으로 정렬되어 있다고 가정할 수 있으며, 그러지 않으면 동작이 정의되지 않습니다. 더 공식적으로는 indices(result)의 모든 i1 < i2에서 full_start_index(i1) <= full_start_index(i2)입니다.

unique_indicestrue이면 구현은 분산되는 모든 result_index 색인이 고유하다고 가정할 수 있습니다. unique_indicestrue이지만 분산되는 색인이 고유하지 않으면 동작이 정의되지 않습니다.

입력

라벨 이름 유형 제약조건
(I1) inputs 텐서의 가변 수 또는 텐서당 양자화 텐서 (C1), (C2), (C4-C6), (C11), (C13), (C18), (C21), (C23-C24)
(I2) scatter_indices 정수 유형의 텐서 (C4), (C15), (C19), (C22)
(I3) updates 텐서의 가변 수 또는 텐서당 양자화 텐서 (C3-C6), (C8)
(I4) update_window_dims si64 유형의 1차원 텐서 상수 (C2), (C4), (C7-C8)
(I5) inserted_window_dims si64 유형의 1차원 텐서 상수 (C2), (C4), (C9-C11)
(I6) input_batching_dims si64 유형의 1차원 텐서 상수 (C2), (C4), (C9), (C12-13), (C17-18), (C20)
(I7) scatter_indices_batching_dims si64 유형의 1차원 텐서 상수 (C14-C18)
(I8) scatter_dims_to_operand_dims si64 유형의 1차원 텐서 상수 (C19-C21)
(I9) index_vector_dim si64 유형의 상수 (C4), (C16), (C19), (C22)
(I10) indices_are_sorted i1 유형의 상수
(I11) unique_indices i1 유형의 상수
(I12) update_computation 기능 (C23)

출력

이름 유형 제약조건
results 텐서의 가변 수 또는 텐서당 양자화 텐서 (C24-C25)

제약조건

  • (C1) same(shape(inputs...)).
  • (C2) `rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims)
    • size(input_batching_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_dimsinput_batching_dims에 해당하는 inputs[0]의 크기 크기는 포함되지 않습니다.
    • combineupdate_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(concatenate(inserted_window_dims, input_batching_dims))
  • (C10) is_sorted(inserted_window_dims)
  • (C11) 0 <= inserted_window_dims < rank(inputs[0]).
  • (C12) is_sorted(input_batching_dims).
  • (C13) 0 <= input_batching_dims < rank(inputs[0])).
  • (C14) is_unique(scatter_indices_batching_dims)입니다.
  • (C15) 0 <= scatter_indices_batching_dims < rank(scatter_indices)입니다.
  • (C16) index_vector_dim not in scatter_indices_batching_dims입니다.
  • (C17) size(input_batching_dims) == size(scatter_indices_batching_dims)입니다.
  • (C18) dim(inputs[0], input_batching_dims...) = dim(scatter_indices, scatter_indices_batching_dims...)입니다.
  • (C19) size(scatter_dims_to_operand_dims) = index_vector_dim < rank(scatter_indices) ? dim(scatter_indices, index_vector_dim) : 1입니다.
  • (C20) is_unique(concatenate(scatter_dims_to_operand_dims, input_batching_dims))
  • (C21) 0 <= scatter_dims_to_operand_dims < rank(inputs[0]).
  • (C22) 0 <= index_vector_dim <= rank(scatter_indices)
  • (C23) update_computation(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>) 유형(여기서 is_promotable(element_type(inputs[i]), Ei))입니다.
  • (C24) shape(inputs...) = shape(results...)입니다.
  • (C25) [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]]
//          ],
//          [
//           [[25, 26], [27, 28], [29, 30], [31, 32]],
//           [[33, 34], [35, 36], [37, 38], [39, 40]],
//           [[41, 42], [43, 44], [45, 46], [47, 48]]
//          ]
//         ]
// %scatter_indices: [
//                    [
//                     [[0, 0], [1, 0], [2, 1]],
//                     [[0, 1], [1, 1], [0, 9]]
//                    ],
//                    [
//                     [[0, 0], [2, 1], [2, 2]],
//                     [[1, 2], [0, 1], [1, 0]]
//                    ]
//                   ]
// %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 = [3, 4],
    inserted_window_dims = [1],
    input_batching_dims = [0],
    scatter_indices_batching_dims = [1],
    scatter_dims_to_operand_dims = [2, 1],
    index_vector_dim = 3>,
  indices_are_sorted = false,
  unique_indices = false
} : (tensor<2x3x4x2xi64>, tensor<2x2x3x2xi64>, tensor<2x2x3x2x2xi64>) -> tensor<2x3x4x2xi64>
// %result: [
//           [
//            [[3, 4], [6, 7], [6, 7], [7, 8]],
//            [[9, 10],[11, 12], [15, 16], [17, 18]],
//            [[17, 18], [19, 20], [22, 23], [24, 25]]
//           ],
//           [
//            [[25, 26], [28, 29], [30, 31], [31, 32]],
//            [[35, 36], [38, 39], [38, 39], [39, 40]],
//            [[41, 42], [44, 45], [46, 47], [47, 48]]
//           ]
//          ]

예시 더보기

select

시맨틱

pred의 상응하는 요소 값에 따라 on_true 또는 on_false 텐서에서 각 요소가 선택되는 result 텐서를 생성합니다. 더 공식적으로는 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 텐서를 생성합니다.

다음 다이어그램은 구체적인 예를 사용하여 result의 요소가 operandsource에서 계산되는 방식을 보여줍니다.

select_and_scatter

더 공식적으로:

  • selected_values = reduce_window_without_init(...)를 다음 입력으로 바꿉니다.

    • inputs = [operand].
    • 있는 그대로 사용되는 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_initreduce_window와 정확히 동일하게 작동합니다. 단, 기본 reduceschedule (축소 참고)에는 init 값이 포함되지 않습니다. 해당 창에 값이 없으면 어떻게 되는지 현재 지정되어 있지 않습니다. (#731)

  • result[result_index] = reduce([source_values], [init_value], [0], scatter) 각 항목의 의미는 다음과 같습니다.

    • source_values = [source[source_index] for source_index in source_indices].
    • selected_values[source_index]operand_indexoperand 요소가 있는 경우 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(tensor<E>, tensor<E>) -> tensor<i1> 유형(여기서 E = element_type(operand))입니다.
  • (C10) scatter(tensor<E>, tensor<E>) -> tensor<E> 유형(여기서 is_promotable(element_type(operand), E))입니다.
  • (C11) shape(operand) = shape(result).
  • (C12) element_type(result) = E.

// %operand: [[1, 5], [2, 5], [3, 6], [4, 4]]
// %source: [[5, 6], [7, 8]]
// %init_value: 0
%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    "stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  window_dimensions = array<i64: 3, 1>,
  window_strides = array<i64: 2, 1>,
  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_transfertrue이면 작업을 통해 데이터를 호스트로 전송합니다. 그렇지 않으면 데이터를 다른 기기로 전송합니다. 이는 구현에 의해 정의됩니다. 이 플래그는 channel_type에 제공된 정보를 복제하므로 향후에는 이 중 하나만 유지할 계획입니다. (#666)

입력

라벨 이름 유형 제약조건
(I1) inputs 텐서 또는 양자화 텐서의 가변 수
(I2) token token
(I3) channel_id si64 유형의 상수
(I4) channel_type DEVICE_TO_DEVICEDEVICE_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 = array<i64: 1, 2>,
  limit_indices = array<i64: 3, 4>,
  strides = array<i64: 1, 1>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>
// % result: [
//            [1, 1],
//            [1, 1]
//           ]

예시 더보기

sort

시맨틱

comparator에 따라 차원 dimension를 따라 inputs의 1차원 슬라이스를 함께 정렬하고 results를 생성합니다.

다른 연산의 유사한 입력과 달리 dimension는 아래에 설명된 시맨틱스를 사용하여 음수 값을 허용합니다. 앞으로는 일관성 때문에 이 방식이 허용되지 않을 수 있습니다. (#1377)

is_stable이 true이면 정렬이 안정적입니다. 즉, 비교 연산자와 동일한 것으로 간주되는 요소의 상대적 순서가 유지됩니다. 단일 입력이 있는 경우 두 요소 e1e2comparator(e1, e2) = comparator(e2, e1) = false인 경우에만 비교 연산자와 동일한 것으로 간주됩니다. 이것이 여러 입력으로 일반화되는 방식은 아래 공식을 참조하세요.

더 공식적으로 index_space(results[0])의 모든 result_index에 대해 다음을 실행합니다.

  • adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension.
  • result_slice = [ri0, ..., :, ..., riR-1], 여기서 riNresult_index에 있는 개별 요소이고 :adjusted_dimension에 삽입됩니다.
  • inputs_together = (inputs[0]..., ..., inputs[N-1]...).
  • results_together[result_slice] = sort(inputs_together[result_slice], comparator_together).
  • 여기서 sort는 왼쪽 인수가 오른쪽 초 인수보다 작으면 comparator_togethertrue를 반환한다고 예상하고 내림차순이 아닌 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]]

예시 더보기

빼기

시맨틱

두 텐서 lhsrhs의 요소별 뺄셈을 수행하고 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) permutationrange(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 = array<i64: 2, 1, 0>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
// %result: [
//           [[1,7], [3,9], [5,11]],
//           [[2,8], [4,10], [6,12]]
//          ]

예시 더보기

triangular_solve

시맨틱

하위 또는 상위 삼각 계수 행렬이 있는 일차 방정식 체계 배치를 풉니다.

더 공식적으로 ab가 주어지면 result[i0, ..., iR-3, :, :]left_sidetrue일 때 op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :] 또는 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 변수를 해결하며 다음 중 하나일 수 있습니다.left_sidefalse

  • NO_TRANSPOSE: a를 있는 그대로 사용하여 작업을 실행합니다.
  • TRANSPOSE: a의 조치 시 작업을 실행합니다.
  • ADJOINT: a의 켤레 전치에 대한 연산을 실행합니다.

입력 데이터는 lowertrue이거나 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) resultEi = 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

시맨틱

두 텐서 lhsrhs에 대해 요소별 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 프로그램에는 StableHLO에 의해 정의되지 않은 작업이 포함되는 경우가 있습니다.

모듈, 함수, 호출, 반환

StableHLO는 ModuleOp, FuncOp, CallOp, ReturnOp에 업스트림 MLIR 작업을 사용합니다. 이는 기존 MLIR 기계와의 상호 운용성을 높이기 위해 수행되었습니다. 많은 유용한 패스가 FuncOp 및 ModuleOp를 대상으로 작성되었고 많은 컴파일 파이프라인에서 이러한 작업이 있을 것으로 예상하기 때문입니다. 이러한 작업에는 완전한 호환성 보장이 적용됩니다. 호환되지 않는 방식 (예: 삭제)으로 이러한 연산이 변경되면 호환성을 유지하기 위해 StableHLO 등가 항목이 추가됩니다.

클로

CHLO 연산에는 StableHLO로 분해하는 상위 수준 작업이 포함됩니다. 현재 CHLO는 호환성이 보장되지 않습니다. 호환성 보장을 위해 직렬화 전에 chlo-legalize-to-stablehlo 통과를 사용해야 합니다.

셰이프 작업

동적 StableHLO 프로그램에서 핵심 MLIR 언어의 특정 작업을 사용하여 셰이프 계산을 수행하는 것은 커뮤니티에서 일반적인 사용 사례입니다. 이러한 작업에는 shape_of 또는 num_elements와 같은 shape 언어 작업, dim 또는 from_elements와 같은 tensor 언어 작업, 기본 제공 index 유형이 포함됩니다.

Dynamism RFC > O2는 이를 범위 외로 나타내지만 index 유형에 관한 일부 지원은 상호 운용을 위해 포함되어 있습니다. 이러한 작업 또는 유형에는 호환성이 보장되지 않습니다. shape-legalize-to-stablehlo 패스를 사용하면 이러한 작업을 완전히 지원되는 StableHLO 작업으로 변환할 수 있습니다.

지원 중단된 작업

MHLO에서 상속된 여러 StableHLO 작업이 있으며, 이러한 작업은 지원 중단되었으며 StableHLO에서 벗어나는 중입니다. 이러한 삭제에 관한 전체 세부정보는 StableHLO v1.0 정리 #2283에서 확인할 수 있습니다. 이러한 지원 중단에 해당하는 추적기 문제는 #2340입니다.

이러한 작업은 몇 가지 카테고리로 분류됩니다.

  • StableHLO 작업의 'HLO에 없음' 카테고리 - 처음에는 StableHLO 연산의 일부였지만 나중에는 적합하지 않은 것으로 간주되었습니다. broadcast, create_token, cross-replica-sum, dot, einsum, torch_index_select, unary_einsum(#3).
  • 미사용 작업 - 이러한 작업은 특정 시점에는 유용했을 수 있지만 작업이 미흡했거나 이러한 작업을 사용하는 파이프라인이 더 이상 필요하지 않도록 리팩터링되었습니다. 여기에는 map, tuple (#598), get_tuple_element, rng, complex 비교 #560, 컨볼루션 window_reversal (#1181)가 포함됩니다.

이러한 작업 중 일부는 기존 작업 (broadcast, create_token, cross-replica-sum, dot, unary_einsum)을 사용하여 표현할 수 있으므로 쉽게 삭제할 수 있으며 기존 호환성 기간 (6개월)이 지나면 삭제됩니다. 다른 작업은 아직 삭제 검토 중입니다 (einsum, get_tuple_element, map, rng torch_index_select, tuple, complex 비교, window_reversal). 커뮤니티 의견을 기다리는 동안 이러한 작업은 삭제되거나 완전한 지원과 함께 사양에 추가됩니다. 이러한 작업 future가 알려질 때까지는 6개월의 호환성만 보장됩니다.

실행

순차적 실행

StableHLO 프로그램은 main 함수에 입력 값을 제공하고 출력 값을 계산하여 실행됩니다. 함수의 출력 값은 상응하는 return 작업에 루팅된 작업의 그래프를 실행하여 계산됩니다.

실행 순서는 Dataflow와 일치한다면(즉, 작업이 사용 전에 실행되는 경우) 구현으로 정의됩니다. StableHLO에서 모든 부작용 작업은 하나의 토큰을 소비하고 하나의 토큰을 생성하므로 (after_all를 통해 여러 토큰을 하나의 토큰으로 멀티플렉스할 수 있음) 부작용의 실행 순서도 Dataflow와 일치합니다. 예를 들어 아래 프로그램에는 %0%1%2return%1%0%2return의 두 가지 실행 순서가 있을 수 있습니다.

func.func @main() -> tensor<f64> {
  %0 = stablehlo.constant dense<1.0> : tensor<f64>
  %1 = stablehlo.constant dense<2.0> : tensor<f64>
  %2 = stablehlo.add %0, %1 : tensor<f64>
  return %2 : tensor<f64>
}

더 공식적으로 StableHLO 프로세스는 1) StableHLO 프로그램, 2) 작업 상태 (아직 실행되지 않음, 이미 실행됨), 3) 프로세스가 작업 중인 중간 값의 조합입니다. 프로세스는 main 함수의 입력 값으로 시작하여, 작업 상태와 중간 값을 업데이트하는 작업의 그래프를 진행하며, 출력 값으로 끝납니다. 추가 형식화는 미정입니다. (#484)

동시 실행

StableHLO 프로그램은 동시에 실행할 수 있으며 둘 다 ui32 유형을 갖는 num_replicas의 2D 프로세스 그리드로 구성됩니다.num_partitions

StableHLO 프로세스 그리드에서 StableHLO 프로세스의 num_replicas * num_partitions가 동시에 실행됩니다. 각 프로세스에는 고유한 process_id = (replica_id, partition_id)가 있습니다. 여기서 replica_ids = range(num_replicas)replica_idpartition_ids = range(num_partitions)partition_id은 둘 다 ui32 유형을 갖습니다.

프로세스 그리드의 크기는 모든 프로그램에서 정적으로 알려져 있으며 (향후 이 프로그램을 StableHLO 프로그램 #650의 명시적인 부분으로 만들 계획임) 프로세스 그리드 내 위치는 모든 프로세스에 대해 통계적으로 알려져 있습니다. 각 프로세스는 replica_idpartition_id 작업을 통해 프로세스 그리드 내 위치에 액세스할 수 있습니다.

프로세스 그리드 내에서는 프로그램이 모두 동일하거나('단일 프로그램, 다중 데이터' 스타일) 다를 수도 있고, '다중 프로그램, 여러 데이터' 스타일로 모두 다를 수도 있고, 그 사이의 프로그램일 수도 있습니다. 향후 GSPMD를 포함하여 병렬 StableHLO 프로그램을 정의하는 다른 관용구를 지원할 계획입니다. (#619)

프로세스 그리드 내에서는 프로세스는 대부분 서로 독립적입니다. 별도의 작업 상태, 별도의 입력/중간/출력 값을 지니며, 아래에 설명된 소수의 집합 작업을 제외하고 대부분의 작업은 프로세스 간에 개별적으로 실행됩니다.

대부분의 연산 실행에는 동일한 프로세스의 값만 사용되므로 일반적으로 이러한 값을 이름으로 참조하는 것이 명확합니다. 하지만 집합 작업의 시맨틱스를 설명하는 경우 이는 충분하지 않으며 특정 프로세스 내에서 name 값을 참조하는 name@process_id 표기법이 발생합니다. (이 관점에서 자격이 없는 namename@(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 x 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_idspartition_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 연산의 해석은 하드웨어 요구사항과 기능에 따라 다를 수 있습니다. 예를 들어 일부 하드웨어는 '역양자화, 부동 소수점 연산 실행, 마지막으로 양자화' 전략을 사용하여 양자화 연산을 해석하도록 선택할 수 있습니다. 나머지는 정수 산술로 전체 계산을 수행할 수 있습니다. 따라서 양자화된 StableHLO 연산의 해석은 전적으로 특정 구현에 의해 결정됩니다. 하이브리드 양자화의 해석(#1575)은 (1792를 통해) 사양에 규정된 대로 시맨틱스를 기반으로 해야 합니다.

오류

StableHLO 프로그램은 개별 작업에 관한 광범위한 제약 조건을 통해 검증되므로 런타임 전에 많은 클래스 오류를 배제합니다. 그러나 오류 조건은 정수 오버플로, 범위를 벗어난 액세스 등을 통해 여전히 발생할 수 있습니다. 명시적으로 호출되지 않는 한 이러한 모든 오류가 구현 정의 동작으로 이어지지만 이는 향후 변경될 수 있습니다. (#1157)

부동 소수점 예외

이 규칙의 예외로, StableHLO 프로그램의 부동 소수점 예외는 동작이 잘 정의되어 있습니다. IEEE-754 표준으로 정의된 예외(잘못된 연산, 0으로 나누기, 오버플로, 언더플로 또는 부정확한 예외)가 발생하는 작업은 표준에서 정의된 대로 기본 결과를 생성하고 해당하는 상태 플래그를 발생시키지 않고 실행을 계속합니다. 이는 표준의 raiseNoFlag 예외 처리와 유사합니다. 비표준 연산 (예: 복잡한 산술 함수 및 특정 초월 함수)의 예외는 구현에 정의되어 있습니다.

모양 불일치

StableHLO는 동적으로 형성된 텐서를 지원합니다. 그러나 도형은 런타임에 동의해야 하며, 그러지 않으면 동작이 정의되지 않습니다. StableHLO는 런타임에 텐서에 특정 형태가 있음을 어설션할 수 있는 작업을 명시적으로 제공하지 않습니다. 올바른 코드를 생성하는 것은 제작자의 책임입니다.

구체적인 예로 아래 프로그램은 유효합니다. 그러나 런타임 시 %arg0%arg1의 정확한 모양은 동일해야 하며, 그렇지 않으면 프로그램의 동작이 정의되지 않습니다.

func.func @foo(%arg0: tensor<?xi32>, %arg1: tensor<?xi32>) -> tensor<?xi32> {
    %0 = stablehlo.add %arg0, %arg1 : tensor<?xi32>
    return %0 : tensor<?xi32>
}

Notation

구문을 설명하기 위해 이 문서에서는 수정된 ISO 버전의 EBNF 구문 (ISO/IEC 14977:1996, 위키백과)을 사용합니다. 두 가지 수정사항이 있습니다. 1) = 대신 ::=를 사용하여 규칙을 정의합니다.

2) , 대신 병치를 사용하여 연결을 표현합니다.

시맨틱을 기술하기 위해(즉, '유형', '상수', '작업' 섹션 내) 배열 작업을 간결하게 표현할 수 있도록 지원되는 확장된 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으로 변환할 수 있는 정확한 형식을 제공하지는 않습니다. 일부 특정 수식이 불분명해 보이는 경우 알려주시면 해당 수식을 개선하도록 노력하겠습니다.

또한 수식은 생략 부호를 사용하여 텐서, 텐서 목록 (예: 다양한 텐서에서 발생할 수 있음)을 비롯한 모든 종류의 목록을 확장합니다. 이는 정확한 공식을 제공하지 않고 (예: 목록이 StableHLO 유형 시스템의 일부도 아님) 이해에 의존하지 않는 또 다른 영역입니다.

C) 마지막으로 주목할 만한 표기법은 암시적 브로드캐스트입니다. StableHLO opset은 암시적 브로드캐스트를 지원하지 않지만 수식은 간결성을 제공하기 위해 지원합니다. 요약하면, 텐서가 예상되는 컨텍스트에서 스칼라를 사용하면 스칼라는 예상 모양으로 브로드캐스트됩니다.

dot_general 예시를 계속 진행하려면 또 다른 제약조건이 0 <= lhs_batching_dimensions < rank(lhs)입니다. dot_general 사양에 정의된 대로 lhs_batching_dimensions는 텐서이지만 0rank(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('유형' 섹션에 정의된 유형), 4) Function('유형' 섹션에 정의된 전역 함수)

컨텍스트에 따라 이름은 다른 값을 참조할 수 있습니다. 더 구체적으로 작업 (및 다른 프로그램 요소에 상응하는 항목)의 '시맨틱' 섹션은 런타임 로직을 정의하므로 모든 입력을 Value로 사용할 수 있습니다. 반면에 작업 (및 이와 동등한 항목)의 '제약 조건' 섹션은 '컴파일 시간' 로직, 즉 일반적으로 런타임 전에 실행되는 로직을 정의합니다. 따라서 상수 입력만 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로 사용할 수 있습니다. 반대로 operandresult는 시맨틱에서는 Value로 사용할 수 있지만 제약 조건에서는 Placeholder로만 사용할 수 있습니다.

함수

유형 생성

유형을 구성하는 데 사용할 수 있는 함수는 없습니다. 대신 유형 문법이 일반적으로 더 간결하기 때문에 직접 사용합니다. 예를 들어 function_type( [tensor_type([], E), tensor_type([], E)], [tensor_type([], E)])가 아닌 (tensor<E>, tensor<E>) -> (tensor<E>)을 사용합니다.

유형에 대한 함수

  • element_type는 텐서 유형과 양자화 텐서 유형에 정의되고 각각 해당하는 TensorType 또는 QuantizedTensorTypeTensorElementType 또는 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) -> Valueis_quantized(x) and quantization_dimension(x) is not None의 바로가기입니다.

  • is_per_tensor_quantized(x: Value | Placeholder | Type) -> Valueis_quantized(x) and quantization_dimension(x) is None의 바로가기입니다.

  • is_promotable(x: Type, y: Type) -> boolx 유형을 y 유형으로 승격할 수 있는지 확인합니다. xyQuantizedTensorElementType이면 프로모션은 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) -> Valueis_quantized_tensor_element_type(x)의 바로가기입니다.

  • is_type_name(x: Value | Placeholder | Type) -> Value. 모든 유형에 사용할 수 있습니다. 예를 들어 xFloatType이면 is_float(x)true를 반환합니다. x가 값 또는 자리표시자인 경우 이 함수는 is_type_name(type(x))의 단축키입니다.

  • max_value(x: Type) -> ValueTensorElementType의 최댓값을 반환합니다. xTensorElementType가 아니면 None를 반환합니다.

  • min_value(x: Type) -> ValueTensorElementType의 가능한 최솟값을 반환합니다. xTensorElementType가 아니면 None를 반환합니다.

  • member_name(x: Value | Placeholder | Type) -> Any. 모든 유형의 모든 구성원 정의 member_name에서 사용할 수 있습니다. 예를 들어 tensor_element_type(x)는 상응하는 TensorTypeTensorElementType 부분을 반환합니다. x가 값 또는 자리표시자인 경우 이 함수는 member_name(type(x))의 단축키입니다. x가 적절한 멤버를 포함하는 유형이 아니거나 이러한 유형의 값 또는 자리표시자가 아니면 None을 반환합니다.

값 구성

  • operation_name(*xs: Value | Type) -> Value. 모든 작업에 사용할 수 있습니다. 예를 들어 add(lhs, rhs)는 2개의 텐서 값 lhsrhs를 사용하고 이러한 입력으로 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의 요소가 색인의 사전 오름차순을 기준으로 오름차순으로 정렬된 경우, 그렇지 않은 경우 false를 반환합니다.true x가 텐서가 아닌 경우 None를 반환합니다.

  • is_unique(x: Value) -> Value는 텐서에 정의되며, x에 중복 요소가 없으면 true를 반환하고 그렇지 않으면 false를 반환합니다. x가 텐서가 아닌 경우 None을 반환합니다.

  • member_name(x: Value) -> Any는 모든 값의 모든 멤버 정의 member_name에 대해 정의됩니다. 예를 들어 real_part(x)는 상응하는 ComplexConstantRealPart 부분을 반환합니다. x가 적절한 멤버가 있는 값이 아닌 경우 None를 반환합니다.

  • same(x: Value) -> Value는 텐서에서 정의되며, x의 요소가 모두 서로 같으면 true를 반환하고 그렇지 않으면 false를 반환합니다. 텐서에 요소가 없으면 '모두 서로 같음'으로 계산됩니다. 즉, 함수는 true를 반환합니다. x가 텐서가 아닌 경우 None를 반환합니다.

  • split(x: Value, num_results: Value, axis: Value) -> Value는 텐서에서 정의되며 축 axis을 따라 xnum_results 슬라이스를 반환합니다. x가 텐서 또는 dim(x, axis) % num_results != 0이 아니면 None를 반환합니다.

  • is_defined_in_parent_scope(x: Value) -> Value는 문자열에 정의되며 x이 관련 작업의 상위 함수와 동일한 범위에서 정의된 함수의 이름인 경우 true를 반환합니다.

  • is_namespaced_op_name(x: Value) -> Value는 문자열에 정의되어 있으며 x이 유효한 작업 이름, 즉 [a-zA-Z][a-zA-Z0-9_]*([.][a-zA-Z0-9_$]+)+ 정규 표현식을 준수하는 경우 true를 반환합니다.

형상 계산

  • axes(x: Value | Placeholder | Type) -> Valuerange(rank(x))의 바로가기입니다.

  • dim(x: Value | Placeholder | Type, axis: Value) -> Valueshape(x)[axis]의 바로가기입니다.

  • dims(x: Value | Placeholder | Type, axes: List) -> Listlist(map(lambda axis: dim(x, axis), axes))의 바로가기입니다.

  • index_space(x: Value | Placeholder | Type) -> Value는 텐서에서 정의되며 사전 오름차순([0, ..., 0], [0, ..., 1], ..., shape(x) - 1)으로 정렬된 해당 TensorType에 대해 size(x) 색인을 반환합니다. x가 텐서 유형, 양자화 텐서 유형, 이러한 유형 중 하나의 값이나 자리표시자가 아니면 None을 반환합니다.

  • rank(x: Value | Placeholder | Type) -> Valuesize(shape(x))의 바로가기입니다.

  • shape(x: Value | Placeholder | Type) -> Valuemember_name를 통해 '유형에 관한 함수' 섹션에 정의되어 있습니다.

  • size(x: Value | Placeholder | Type) -> Valuereduce(lambda x, y: x * y, shape(x))의 바로가기입니다.

양자화 계산

  • def baseline_element_type(x: Value | Placeholder | Type) -> Typeelement_type(baseline_type(x))의 바로가기입니다.

  • baseline_type는 텐서 유형과 양자화 텐서 유형에 정의되며 '기준'으로 변환됩니다. 즉, 형태는 동일하지만 양자화 매개변수가 있는 요소 유형의 양자화 매개변수가 기본값으로 재설정됩니다. 이 방법은 텐서 및 양자화 텐서 유형을 균일하게 비교하는 데 사용되며, 매우 자주 필요합니다. 양자화 유형의 경우 이를 통해 양자화 매개변수를 무시하는 유형을 비교할 수 있습니다. 즉, shape, storage_type, expressed_type, storage_min, storage_max, quantization_dimension (축별 양자화 유형)는 모두 일치해야 하지만 scaleszero 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, result_type: Type) -> Value:
  assert is_float(x) and is_quantized(result_type)
  zero_points = compute_zero_points(result_type, TensorType(shape(x), storage_type(result_type)))
  converted_zero_points = convert(zero_points, expressed_type(result_type))
  converted_min = convert(storage_min(result_type), expressed_type(result_type))
  converted_max = convert(storage_max(result_type), expressed_type(result_type))

  x_scaled = x / compute_scales(result_type, type(x))
  x_scaled_add_zp = x_scaled + converted_zero_points
  x_clamped = clamp(converted_min, x_scaled_add_zp, converted_max)
  x_rounded = round_nearest_even(x_clamped)
  return convert(x_rounded, result_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)
  • hybrid_dequantize_then_op는 부동 소수점의 lhs와 양자화 유형의 rhs를 허용하는 하이브리드 오퍼레이션의 가중치 전용 양자화를 지정하는 데 사용됩니다. 양자화된 입력을 표현된 유형으로 역양자화하고 부동 소수점 형식으로 계산을 수행합니다. float lhs 텐서의 요소 유형과 양자화 rhs 텐서의 표현 유형은 동일해야 합니다.
def hybrid_dequantize_then_op(op, lhs, rhs):
  assert(is_float(lhs) and is_quantized(rhs) and element_type(lhs) == expressed_type(rhs))
  return op(lhs, dequantize(rhs))

그리드 계산

  • 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' 섹션을 참조하세요.

역동적인

StableHLO 값에는 동적 크기(예: tensor<?xi64>)가 포함될 수 있습니다. 하지만 StableHLO 값은 동적 측정기준 수 (순위가 지정되지 않은 동적, 예: tensor<*xi64>)를 포함할 수 없습니다. 크기에 제약이 있더라도 피연산자와 결과에서 동적 측정기준 크기를 사용할 수 있습니다. 제약 조건은 가능한 경우 정적으로 확인됩니다. 그러지 않으면 제약 조건은 런타임으로 지연되고 불일치로 인해 정의되지 않은 동작이 발생합니다. 아래에서 예를 확인하세요.

단항 요소별 연산의 도형 불일치

다음 장난감 프로그램을 생각해 보세요.

func.func @foo(%arg0: tensor<?xf64>) {
  %0 = stablehlo.abs %arg0 : (tensor<?xf64>) -> tensor<2xf64>
  return
}

이러한 프로그램은 흔하지 않습니다. 결과의 형태는 알고 있지만 입력의 형태는 모르는 것이 일반적이기 때문입니다. 그럼에도 불구하고 이는 유효한 StableHLO 프로그램입니다. 이 프로그램에서는 abs 연산의 정적인 유효성을 정적으로 검사할 수 없습니다. 피연산자의 정확한 형태를 알 수 없기 때문입니다. 그러나 도형은 확실히 호환되며 정적으로 확인할 수 있습니다. 즉, ?가 런타임에 2이 될 수 있으며 문제가 없습니다. 하지만 ?가 다른 정수가 될 수도 있으며, 이 경우 동작이 정의되지 않습니다.

결과에서 크기 크기가 동적인 경우 정의되지 않은 동작이 있을 수 없습니다. 실제로 '예상되는' 크기가 없으므로 불일치가 발생할 수 없습니다.

바이너리 요소별 연산의 도형 불일치

다음 장난감 프로그램을 생각해 보세요.

func.func @foo(%arg0: tensor<?xf64>, %arg1: tensor<?xf64>) {
  %0 = stablehlo.add %arg0, %arg0 : (tensor<?xf64>, tensor<?xf64>) -> tensor<?xf64>
  return
}

바이너리 요소별 연산의 경우 입력 모양과 결과가 런타임에 일치해야 합니다. 컴파일 시점에 정적 차원은 동일해야 하며, 그렇지 않으면 호환되기만 하면 됩니다. 입력에서 하나의 측정기준이 동적인 경우 동적 크기가 다른 피연산자의 해당 크기와 일치하지 않을 수 있으므로 (정적 또는 동적) 런타임 시 정의되지 않은 동작이 발생할 수 있습니다. 모든 입력이 정적인 경우 결과가 동적인지 여부는 중요하지 않습니다. 정적으로 알려진 측정기준은 정적으로 검사되며 동적 측정기준은 제약을 적용하지 않습니다.

출력 형태를 피연산자로 사용하는 작업의 형태 불일치

다음 장난감 프로그램을 생각해 보세요.

func.func @foo(%arg0: tensor<2xi32>) {
  %0 = stablehlo.dynamic_iota %arg0, dim = 0 : (tensor<2xi32>) -> tensor<3x4xi64>
  return
}

런타임 시 셰이프 피연산자의 값은 결과의 셰이프와 일치해야 하며, 그렇지 않으면 동작이 정의되지 않습니다. 즉, 런타임 시 %arg0의 값은 dense<[3, 4]> : tensor<2xi32>여야 합니다. 셰이프 피연산자가 상수이면 정적으로 확인할 수 있습니다. 결과 모양이 완전히 동적이면 불일치가 발생할 수 없습니다.