Especificação do StableHLO

O StableHLO é um conjunto de operações para operações de alto nível (HLO) em modelos de machine learning (ML). O StableHLO funciona como uma camada de portabilidade entre diferentes frameworks e compiladores de ML: os frameworks que produzem programas StableHLO são compatíveis com os compiladores que consomem esses programas.

Nosso objetivo é simplificar e acelerar o desenvolvimento de ML criando mais interoperabilidade entre vários frameworks de ML (como TensorFlow, JAX e PyTorch) e compiladores de ML (como XLA e IREE). Para isso, este documento fornece uma especificação para a linguagem de programação StableHLO.

Esta especificação contém três seções principais. Primeiro, a seção Programas descreve a estrutura dos programas StableHLO, que consistem em funções StableHLO, que consistem em operações StableHLO. Nessa estrutura, a seção Ops especifica a semântica de operações individuais. A seção Execução fornece semântica para todas essas operações executadas juntas em um programa. Por fim, a seção Notação aborda a notação usada em toda a especificação.

Para conferir a especificação de uma versão anterior do StableHLO, abra o repositório na versão marcada de interesse. Por exemplo, a especificação do StableHLO v0.19.0. Para conferir as mudanças que ocorreram em cada aumento de versão secundária do StableHLO, consulte o registro de versão em VhloDialect.td.

Programas

Program ::= {Func}

Os programas do StableHLO consistem em um número arbitrário de funções do StableHLO. Confira abaixo um exemplo de programa com uma função @main que tem três entradas (%image, %weights e %bias) e uma saída. O corpo da função tem seis operações.

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

Funções

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

As funções StableHLO (também chamadas de funções nomeadas) têm um identificador, entradas/saídas e um corpo. No futuro, planejamos introduzir metadados adicionais para funções e alcançar uma compatibilidade melhor com HLO (#425, #626, #740, #744).

Identificadores

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

Os identificadores StableHLO são semelhantes aos identificadores em muitas linguagens de programação, com duas peculiaridades: 1) todos os identificadores têm sigilos que distinguem diferentes tipos de identificadores; 2) os identificadores de valor podem ser completamente numéricos para simplificar a geração de programas StableHLO.

Tipos

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

Os tipos do StableHLO são categorizados em tipos de valor (também chamados de tipos de primeira classe), que representam valores do StableHLO, e tipos não de valor, que descrevem outros elementos do programa. Os tipos do StableHLO são semelhantes aos de muitas linguagens de programação. A principal peculiaridade é a natureza específica do domínio do StableHLO, que resulta em alguns resultados incomuns (por exemplo, tipos escalares não são tipos de valor).

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

Os tipos de tensor representam tensores, ou seja, matrizes multidimensionais. Eles têm um formato e um tipo de elemento, em que um formato representa tamanhos de dimensão não negativos ou desconhecidos na ordem crescente das dimensões correspondentes (também chamadas de eixos) numeradas de 0 a R-1. O número de dimensões R é chamado de cardinalidade. Por exemplo, tensor<2x3xf32> é um tipo de tensor com formato 2x3 e tipo de elemento f32. Ela tem duas dimensões (ou seja, dois eixos): a 0ª e a 1ª, cujos tamanhos são 2 e 3. A classificação é 2.

As dimensões podem ser parcialmente ou totalmente desconhecidas (dinâmicas). Por exemplo, tensor<?x2xf64> é parcialmente desconhecida e tensor<?x?xf64> é totalmente desconhecida. Os tamanhos de dimensões dinâmicas são representados usando um ?. Não é possível remover a classificação das formas.

No futuro, planejamos estender os tipos de tensor além dos tamanhos de dimensão e tipos de elementos, por exemplo, para incluir layouts (#629) e esparsidade (#1078).

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

Tipos de elementos quantizados representam valores inteiros de um tipo de armazenamento no intervalo de storage_min a storage_max (inclusive) que correspondem a valores de ponto flutuante de um tipo expresso. Para um determinado valor inteiro i, o valor de ponto flutuante correspondente f pode ser calculado como f = (i - zero_point) * scale, em que scale e zero_point são chamados de parâmetros de quantização. Os storage_min e storage_max são opcionais na gramática, mas têm valores padrão de min_value(storage_type) e max_value(storage_type), respectivamente. Os tipos de elementos quantizados têm as seguintes restrições:

• (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) Se is_empty(quantization_dimension), então size(scales) = 1.
• (C11) 0 <= quantization_dimension.

No momento, QuantizationScale é uma constante de ponto flutuante, mas há um grande interesse em escalas baseadas em números inteiros, representadas com multiplicadores e deslocamentos. Estamos planejando explorar isso em breve (#1404).

Há uma discussão em andamento sobre a semântica de QuantizationZeroPoint, incluindo o tipo, os valores e se pode haver apenas um ou potencialmente vários pontos zero em um tipo de tensor quantizado. Com base nos resultados dessa discussão, a especificação sobre pontos zero pode mudar no futuro (#1405).

Outra discussão em andamento envolve a semântica de QuantizationStorageMin e QuantizationStorageMax para determinar se alguma restrição deve ser imposta a esses valores e aos valores de tensores quantizados (#1406).

Por fim, planejamos explorar a representação de escalas desconhecidas e pontos zero, assim como planejamos explorar a representação de tamanhos de dimensões desconhecidas (#1407).

Os tipos de tensores quantizados representam tensores com elementos quantizados. Esses tensores são exatamente iguais aos tensores regulares, exceto que os elementos deles têm tipos de elementos quantizados, em vez de tipos de elementos regulares.

Em tensores quantizados, a quantização pode ser por tensor, ou seja, ter um scale e um zero_point para todo o tensor, ou por eixo, ou seja, ter vários scales e zero_points, um par por fração de uma dimensão específica quantization_dimension. Mais formalmente, em um tensor t com quantização por eixo, há dim(t, quantization_dimension) slices de quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] etc. Todos os elementos no iº slice usam scales[i] e zero_points[i] como parâmetros de quantização. Os tipos de tensores quantizados têm as seguintes restrições:

• Para quantização por tensor:
  • Sem outras restrições.
• Para quantização por eixo:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Os tipos de token representam tokens, ou seja, valores opacos produzidos e consumidos por algumas operações. Os tokens são usados para impor a ordem de execução das operações, conforme descrito na seção Execução.

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

Tipos de buffer representam buffers. Por exemplo, no XLA, os buffers são matrizes multidimensionais com armazenamento consistente. Assim como os tipos de tensor, os tipos de buffer têm uma forma e um tipo de elemento, em que uma forma representa tamanhos de dimensão não negativos ou desconhecidos na ordem crescente das dimensões correspondentes (também chamadas de eixos) numeradas de 0 a R-1. O número de dimensões R é chamado de classificação. Por exemplo, memref<2x3xf32> é um tipo de buffer com formato 2x3 e tipo de elemento f32. Ela tem duas dimensões (ou seja, dois eixos): a dimensão 0 e a dimensão 1, cujos tamanhos são 2 e 3. A classificação é 2.

Os buffers podem ser alocados usando um custom_call para CreateBuffer ou Pin e desalocados por um custom_call para Unpin. Somente operações custom_call podem ler e gravar o conteúdo dentro dos buffers. Consulte custom_call para mais detalhes.

Tipos de tupla representam tuplas, ou seja, listas heterogêneas. As tuplas são um recurso legado que existe apenas para compatibilidade com HLO. Em HLO, as tuplas são usadas para representar entradas e saídas variadas. No StableHLO, entradas e saídas variádicas são compatíveis de forma nativa, e o único uso de tuplas é para representar de forma abrangente a ABI do HLO, em que, por exemplo, T, tuple<T> e tuple<tuple<T>> podem ser materialmente diferentes dependendo de uma implementação específica. No futuro, planejamos fazer mudanças na ABI do HLO, o que pode permitir remover tipos de tupla do StableHLO (#598).

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

Os tipos de elementos representam elementos de tipos de tensor. Ao contrário de muitas linguagens de programação, esses tipos não são de primeira classe no StableHLO. Isso significa que os programas StableHLO não podem representar diretamente valores desses tipos. Como resultado, é idiomático representar valores escalares do tipo T com valores de tensor de dimensão zero do tipo tensor<T>.

• O tipo booleano representa os valores booleanos true e false.
• Tipos de números inteiros podem ser assinados (si) ou não assinados (ui) e têm uma das larguras de bits compatíveis (2, 4, 8, 16, 32 ou 64). Os tipos siN assinados representam valores inteiros de -2^(N-1) a 2^(N-1)-1, inclusive, e os tipos uiN não assinados representam valores inteiros de 0 a 2^N-1, inclusive.
• Os tipos de ponto flutuante podem ser um dos seguintes:
  • Números de ponto flutuante de 8 bits f8E3M4, f8E4M3 e f8E5M2 seguindo as convenções IEEE-754.
  • Tipos f8E4M3FN e f8E5M2 correspondentes, respectivamente, às codificações E4M3 e E5M2 do formato FP8 descritas em Formatos FP8 para aprendizado profundo.
  • Tipos f8E4M3FNUZ e f8E5M2FNUZ correspondentes às codificações E4M3 e E5M2 dos formatos FP8 descritos em Formatos numéricos de 8 bits para redes neurais profundas.
  • Tipo f8E4M3B11FNUZ correspondente à codificação E4M3 dos formatos FP8 descritos em Treinamento e inferência de ponto flutuante híbrido de 8 bits (HFP8) para redes neurais profundas.
  • Tipo bf16 correspondente ao formato bfloat16 descrito em BFloat16: o segredo do alto desempenho em Cloud TPUs.
  • Os tipos f16, f32 e f64 correspondem, respectivamente, aos formatos binary16 ("meia precisão"), binary32 ("precisão simples") e binary64 ("precisão dupla") descritos no padrão IEEE 754.
  • O tipo tf32 corresponde ao formato TensorFloat32 e tem suporte limitado no StableHLO.
  • f4E2M1FN, f6E2M3FN, f6E3M2FN e f8E8M0FNU MX (microescalonamento) descritos na Especificação de formatos de microescalonamento do OCP.
• Tipos complexos representam valores complexos que têm uma parte real e uma parte imaginária do mesmo tipo de elemento. Os tipos complexos compatíveis são complex<f32> (ambas as partes são do tipo f32) e complex<f64> (ambas as partes são do tipo f64).

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

Os tipos de função representam funções nomeadas e anônimas. Elas têm tipos de entrada (a lista de tipos no lado esquerdo de ->) e tipos de saída (a lista de tipos no lado direito de ->). Em muitas linguagens de programação, os tipos de função são de primeira classe, mas não no StableHLO.

StringType ::= 'string'

O tipo de string representa sequências de bytes. Ao contrário de muitas linguagens de programação, o tipo de string não é de primeira classe no StableHLO e é usado apenas para especificar metadados estáticos para elementos de programa.

Operações

As operações StableHLO (também chamadas de ops) representam um conjunto fechado de operações de alto nível em modelos de machine learning. Conforme discutido acima, a sintaxe do StableHLO é muito inspirada no MLIR, que não é necessariamente a alternativa mais ergonômica, mas é provavelmente a mais adequada para o objetivo do StableHLO de criar mais interoperabilidade entre frameworks e compiladores de ML.

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

As operações do StableHLO (também chamadas de ops) têm um nome, entradas/saídas e uma assinatura. O nome consiste no prefixo stablehlo. e um mnemônico que identifica exclusivamente uma das operações compatíveis. Confira abaixo uma lista completa de todas as operações compatíveis.

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

As operações consomem entradas e produzem saídas. As entradas são categorizadas em valores de entrada (calculados durante a execução), funções de entrada (fornecidas estaticamente, porque no StableHLO as funções não são valores de primeira classe) e atributos de entrada (também fornecidos estaticamente). O tipo de entradas e saídas consumidas e produzidas por uma operação depende do mnemônico dela. Por exemplo, a operação add consome dois valores de entrada e produz um valor de saída. Em comparação, a operação select_and_scatter consome três valores de entrada, duas funções de entrada e três atributos de entrada.

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

As funções de entrada (também chamadas de funções anônimas) são muito parecidas com as funções nomeadas, exceto que: 1) elas não têm um identificador (daí o nome "anônimas"); 2) elas não declaram tipos de saída (os tipos de saída são inferidos da operação return na função).

A sintaxe das funções de entrada inclui uma parte não usada no momento (consulte a produção Unused acima), que está lá para compatibilidade com MLIR. No MLIR, existe um conceito mais geral de "regiões", que podem ter vários "blocos" de operações conectados por operações de salto. Esses blocos têm IDs que correspondem à produção Unused, para que possam ser diferenciados. O StableHLO não tem operações de salto, então a parte correspondente da sintaxe MLIR não é usada, mas ainda está lá.

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

Os atributos de entrada têm um nome e um valor que é uma das constantes compatíveis. Eles são a principal maneira de especificar metadados estáticos para elementos de programa. Por exemplo, a operação concatenate usa o atributo dimension para especificar a dimensão ao longo da qual os valores de entrada são concatenados. Da mesma forma, a operação slice usa vários atributos, como start_indices e limit_indices, para especificar os limites usados para dividir o valor de entrada.

No momento, os programas StableHLO em uso às vezes contêm atributos que não são descritos neste documento. No futuro, planejamos incorporar esses atributos ao conjunto de operações do StableHLO ou proibir que eles apareçam em programas do StableHLO. Enquanto isso, confira a lista de atributos:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Metadados de local (#594).

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

A assinatura da operação consiste nos tipos de todos os valores de entrada (a lista de tipos no lado esquerdo de ->) e nos tipos de todos os valores de saída (a lista de tipos no lado direito de ->). Falando estritamente, os tipos de entrada são redundantes, e os tipos de saída também são quase sempre redundantes, porque, para a maioria das operações do StableHLO, os tipos de saída podem ser inferidos das entradas. No entanto, a assinatura da operação faz parte da sintaxe do StableHLO para compatibilidade com MLIR.

Confira abaixo um exemplo de operação cujo mnemônico é select_and_scatter. Ele consome três valores de entrada (%operand, %source e %init_value), duas funções de entrada e três atributos de entrada (window_dimensions, window_strides e padding). Observe como a assinatura da operação inclui apenas os tipos dos valores de entrada, mas não os tipos de funções e atributos de entrada, que são fornecidos inline.

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

Constantes

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

As constantes do StableHLO têm um literal e um tipo que representam juntos um valor do StableHLO. Em geral, o tipo faz parte da sintaxe constante, exceto quando não há ambiguidade. Por exemplo, uma constante booleana tem o tipo i1, enquanto uma constante de número inteiro pode ter vários tipos possíveis.

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

As constantes booleanas representam os valores booleanos true e false. As constantes booleanas têm o tipo i1.

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

Constantes de números inteiros representam valores inteiros usando strings que usam notação decimal ou hexadecimal. Outras bases, como binária ou octal, não são compatíveis. As constantes inteiras têm as seguintes restrições:

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

As constantes de ponto flutuante representam valores de ponto flutuante usando strings que usam notação decimal ou científica. Além disso, a notação hexadecimal pode ser usada para especificar diretamente os bits subjacentes no formato de ponto flutuante do tipo correspondente. As constantes de ponto flutuante têm as seguintes restrições:

• (C1) Se uma notação não hexadecimal for usada, is_wellformed(float_literal, float_type).
• (C2) Se a notação hexadecimal for usada, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Constantes complexas representam valores complexos usando listas de uma parte real (primeiro) e uma parte imaginária (segundo). Por exemplo, (1.0, 0.0) : complex<f32> representa 1.0 + 0.0i e (0.0, 1.0) : complex<f32> representa 0.0 + 1.0i. A ordem em que essas partes são armazenadas na memória é definida pela implementação. As constantes complexas têm as seguintes restrições:

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

As constantes de tensor representam valores de tensor usando listas aninhadas especificadas pela notação do NumPy. Por exemplo, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> representa um valor de tensor com o seguinte mapeamento de índices para elementos: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. A ordem em que esses elementos são armazenados na memória é definida pela implementação. As constantes de tensor têm as seguintes restrições:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), em que:
  • 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)), em que:
  • 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:]).
  • Caso contrário, false.

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

As constantes de tensor quantizado representam valores de tensor quantizado usando a mesma notação das constantes de tensor, com elementos especificados como constantes do tipo de armazenamento. As constantes de tensor quantizadas têm as seguintes restrições:

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

Os literais de string consistem em bytes especificados usando caracteres ASCII e sequências de escape. Eles são independentes de codificação, então a interpretação desses bytes é definida pela implementação. Os literais de string têm o tipo string.

Operações

abs

Semântica

Executa uma operação abs elemento a elemento no tensor operand e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para números inteiros com sinal: módulo inteiro.
• Para pontos flutuantes: abs do IEEE-754.
• Para números complexos: módulo complexo.
• Para tipos quantizados: dequantize_op_quantize(abs, operand, type(result)).

Entradas

Saídas

Restrições

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) é definido como:
  • complex_element_type(element_type(operand)) se is_complex(operand).
  • baseline_element_type(operand) caso contrário.

Exemplos

// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand)< : (t>ens>or3xi32<) - t>ensor3xi32
// %result: [2, 0, 2]

 Mais exemplos

adicionar

Semântica

Realiza a adição de elementos de dois tensores lhs e rhs e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para booleanos: OR lógico.
• Para números inteiros: adição de números inteiros.
• Para pontos flutuantes: addition do IEEE-754.
• Para números complexos: adição complexa.
• Para tipos quantizados: dequantize_op_quantize(add, lhs, rhs, type(result)).

Entradas

Saídas

Restrições

• Se a operação usar tensores não quantizados:
  • (C1) type(lhs) = type(rhs) = type(result).
• Se a operação usar tensores quantizados:
  • (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) Se is_per_axis_quantized(lhs), então quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Se is_per_axis_quantized(rhs), então quantization_dimension(rhs) = quantization_dimension(result).

Exemplos

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

 Mais exemplos

after_all

Semântica

Garante que as operações que produzem o inputs sejam executadas antes de qualquer operação que dependa de result. A execução dessa operação não faz nada. Ela só existe para estabelecer dependências de dados de result para inputs.

Entradas

Saídas

Exemplos

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

 Mais exemplos

all_gather

Semântica

Em cada grupo de processos na grade de processos do StableHLO, concatena os valores dos tensores operands de cada processo ao longo de all_gather_dim e produz tensores results.

A operação divide a grade de processo do StableHLO em process_groups, que é definida da seguinte forma:

• cross_replica(replica_groups)se channel_id <= 0 and use_global_device_ids = false.
• cross_replica_and_partition(replica_groups)se channel_id > 0 and use_global_device_ids = false.
• flattened_ids(replica_groups)se channel_id > 0 and use_global_device_ids = true.

Depois, em cada process_group:

• operands...@receiver = [operand@sender for sender in process_group] para todos os receiver em process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) para todos os process em process_group.

Entradas

Saídas

Restrições

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) é definido como:
  • num_replicas se cross_replica for usado.
  • num_replicas se cross_replica_and_partition for usado.
  • num_processes se flattened_ids for usado.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Se use_global_device_ids = true, então channel_id > 0.
• (C6) type(results...) = type(operands...), exceto:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Exemplos

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

 Mais exemplos

all_reduce

Semântica

Em cada grupo de processos na grade de processos do StableHLO, aplica uma função de redução computation aos valores dos tensores operands de cada processo e produz tensores results.

A operação divide a grade de processo do StableHLO em process_groups, que é definida da seguinte forma:

• cross_replica(replica_groups)se channel_id <= 0 and use_global_device_ids = false.
• cross_replica_and_partition(replica_groups)se channel_id > 0 and use_global_device_ids = false.
• flattened_ids(replica_groups)se channel_id > 0 and use_global_device_ids = true.

Depois, em cada process_group:

• results...@process[result_index] = exec(schedule) para alguma árvore binária schedule em que:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule é uma árvore binária definida pela implementação cuja travessia em ordem é to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Entradas

Saídas

Restrições

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) é definido como:
  • num_replicas se cross_replica for usado.
  • num_replicas se cross_replica_and_partition for usado.
  • num_processes se flattened_ids for usado.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Se use_global_device_ids = true, então channel_id > 0.
• (C5) computation tem o tipo (tensor<E>, tensor<E>) -> (tensor<E>) em que is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Exemplos

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

 Mais exemplos

all_to_all

Semântica

Em cada grupo de processos na grade de processos do StableHLO, divide os valores dos tensores operands ao longo de split_dimension em partes, dispersa as partes divididas entre os processos, concatena as partes dispersas ao longo de concat_dimension e produz tensores results. A operação divide a grade de processo do StableHLO em process_groups, que é definida da seguinte forma:

• cross_replica(replica_groups) se channel_id <= 0.
• cross_partition(replica_groups) se channel_id > 0.

Depois, em cada process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) para todos os sender em process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] em que receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Entradas

Saídas

Restrições

• (C1) 0 <= split_dimension < rank(operands...).
• (C2) dim(operands..., split_dimension) % split_count = 0.
• (C3) 0 <= concat_dimension < rank(operands...).
• (C4) 0 < split_count.
• (C5) is_unique(replica_groups).
• (C6) size(replica_groups) é definido como:
  • num_replicas se cross_replica for usado.
  • num_partitions se cross_partition for usado.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...), exceto se split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Exemplos

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

 Mais exemplos

e

Semântica

Executa a operação AND bit a bit de dois tensores lhs e rhs e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para booleanos: AND lógico.
• Para números inteiros: AND bit a bit.

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

atan2

Semântica

Executa a operação atan2 com elementos nos tensores lhs e rhs e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para pontos flutuantes: atan2 do IEEE-754.
• Para números complexos: atan2 complexo.
• Para tipos quantizados: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

// %lhs: [0.0, 1.0, -1.0]
// %rhs: [0.0, 0.0, 0.0]
%result = "stablehlo.atan2"(%lhs, %rhs)< : (t>ensor3xf<64, t>ens>or3xf64<) - t>ensor3xf64
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

 Mais exemplos

batch_norm_grad

Semântica

Calcula gradientes de várias entradas de batch_norm_training usando backpropagation de grad_output e produz tensores grad_operand, grad_scale e grad_offset. De maneira mais formal, essa operação pode ser expressa como uma decomposição em operações StableHLO atuais usando a sintaxe do Python da seguinte maneira:

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

Para tipos quantizados, realiza 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)).

Entradas

Saídas

Restrições

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale e grad_offset têm o mesmo baseline_element_type.
• (C3) operand, grad_output e grad_operand têm o mesmo formato.
• (C4) scale, mean, variance, grad_scale e grad_offset têm a mesma forma.
• (C5) size(scale) = dim(operand, feature_index).

Exemplos

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
// %grad_output: [
//                [[0.1, 0.1], [0.1, 0.1]],
//                [[0.1, 0.1], [0.1, 0.1]]
//               ]
%grad_operand, %grad_scale, %grad_offset =
"stablehlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf64,
 <    tenso>r2x>2x2xf64)< - (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %grad_operand: [
//                 [[0.0, 0.0], [0.0, 0.0]],
//                 [[0.0, 0.0], [0.0, 0.0]]
//                ]
// %grad_scale:  [0.0, 0.0]
// %grad_offset: [0.4, 0.4]

batch_norm_inference

Semântica

Normaliza o tensor operand em todas as dimensões, exceto a feature_index, e produz um tensor result. De maneira mais formal, essa operação pode ser expressa como uma decomposição em operações StableHLO atuais usando a sintaxe do Python da seguinte maneira:

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)

Para tipos quantizados, realiza 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)).

Entradas

Saídas

Restrições

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance e result têm o mesmo baseline_element_type.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index).
• (C5) size(mean) = dim(operand, feature_index).
• (C6) size(variance) = dim(operand, feature_index).
• (C7) baseline_type(operand) = baseline_type(result).

Exemplos

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
%result = "stablehlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ens>or2xf64<) - tenso>r2x2x2xf64
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

batch_norm_training

Semântica

Calcula a média e a variância em todas as dimensões, exceto a feature_index, e normaliza o tensor operand, produzindo os tensores output, batch_mean e batch_var. De maneira mais formal, essa operação pode ser expressa como uma decomposição em operações StableHLO atuais usando a sintaxe do Python da seguinte maneira:

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

Para tipos quantizados, realiza 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)).

Entradas

Saídas

Restrições

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var e output têm o mesmo baseline_element_type.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index).
• (C5) size(batch_mean) = dim(operand, feature_index).
• (C6) size(batch_var) = dim(operand, feature_index).
• (C7) baseline_type(output) = baseline_type(operand).

Exemplos

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
%output, %batch_mean, %batch_var = "stablehlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ens>or2xf64) -
 <   (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %output: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]
// %batch_mean: [2.0, 3.0]
// %batch_var: [1.0, 1.0]

bitcast_convert

Semântica

Realiza uma operação bitcast no tensor operand e produz um tensor result em que os bits de todo o tensor operand são reinterpretados usando o tipo do tensor result.

Mais formalmente, considerando E = element_type(operand), E' = element_type(result) e R = rank(operand):

• Se num_bits(E') < num_bits(E), bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1]).
• Se num_bits(E') > num_bits(E), bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :]).
• Se num_bits(E') = num_bits(E), bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1]).

bits retorna a representação na memória de um determinado valor, e o comportamento dele é definido pela implementação porque a representação exata de tensores e tipos de elementos também é definida pela implementação.

Entradas

Saídas

Restrições

• (C1) Considerando E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) e R = rank(operand):
  • Se num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Se num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) para todos os 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Se num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) para todos os 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Se is_complex(operand) or is_complex(result), então is_complex(operand) and is_complex(result).

Exemplos

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

 Mais exemplos

broadcast_in_dim

Semântica

Expande as dimensões e/ou a classificação de um tensor de entrada duplicando os dados no tensor operand e produz um tensor result. De maneira mais formal, result[result_index] = operand[operand_index], em que para todo d em axes(operand):

• operand_index[d] = 0 se dim(operand, d) = 1.
• operand_index[d] = result_index[broadcast_dimensions[d]] caso contrário.

Entradas

Saídas

Restrições

• (C1) element_type(result) é dado por:
  • element_type(operand), se !is_per_axis_quantized(operand).
  • element_type(operand), exceto que quantization_dimension(operand), scales(operand) e zero_points(operand) podem ser diferentes de quantization_dimension(result), scales(result) e zero_points(result), respectivamente, caso contrário.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Para todos os d em axes(operand):
  • dim(operand, d) = 1 ou
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Se is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Se dim(operand, quantization_dimension(operand)) = 1, então scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Exemplos

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

 Mais exemplos

caso

Semântica

Produz a saída da execução de exatamente uma função de branches, dependendo do valor de index. De maneira mais formal, result = selected_branch() em que:

• selected_branch = branches[index] se 0 <= index < size(branches).
• selected_branch = branches[-1] caso contrário.

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

cbrt

Semântica

Executa a operação de raiz cúbica elemento a elemento no tensor operand e produz um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para pontos flutuantes: rootn(x, 3) do IEEE-754.
• Para números complexos: raiz cúbica complexa.
• Para tipos quantizados: dequantize_op_quantize(cbrt, operand, type(result))

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

ceil

Semântica

Realiza o teto de elemento a elemento do tensor operand e produz um tensor result. Implementa a operação roundToIntegralTowardPositive da especificação IEEE-754. Para tipos quantizados, realiza dequantize_op_quantize(ceil, operand, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

cholesky

Semântica

Calcula a decomposição de Cholesky de um lote de matrizes.

Mais formalmente, para todo i em index_space(result), result[i0, ..., iR-3, :, :] é uma decomposição de Cholesky de a[i0, ..., iR-3, :, :], na forma de uma matriz triangular inferior (se lower for true) ou triangular superior (se lower for false). Os valores de saída no triângulo oposto, ou seja, o triângulo estrito superior ou inferior, respectivamente, são definidos pela implementação.

Se houver i em que a matriz de entrada não seja uma matriz hermitiana positiva definida, o comportamento será indefinido.

Para tipos quantizados, realiza dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

limitar

Semântica

Fixa cada elemento do tensor operand entre um valor mínimo e máximo e produz um tensor result. Mais formalmente, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), em que min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Para tipos quantizados, realiza dequantize_op_quantize(clamp, min, operand, max, type(result)).

Impor uma ordenação em números complexos envolve semântica surpreendente. Por isso, no futuro, planejamos remover o suporte a números complexos para essa operação (#560).

Entradas

Saídas

Restrições

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

Exemplos

// %min: [5, 10, 15]
// %operand: [3, 13, 23]
// %max: [10, 15, 20]
%result = "stablehlo.clamp"(%min, %operand, %max)< : (t>ensor3xi<32, t>ensor3xi<32, t>ens>or3xi32<) - t>ensor3xi32
// %result: [5, 13, 20]

 Mais exemplos

collective_broadcast

Semântica

Em cada grupo de processos na grade de processos do StableHLO, envie o valor do tensor operand do processo de origem para os processos de destino e produza um tensor result.

A operação divide a grade de processo do StableHLO em process_groups, que é definida da seguinte forma:

• cross_replica(replica_groups) se channel_id <= 0.
• cross_partition(replica_groups) se channel_id > 0.

Depois disso, result@process é dado por:

• operand@process_groups[i, 0] se existir um i para que o processo esteja em process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) caso contrário.

Entradas

Saídas

Restrições

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N em que N é definido como:
  • num_replicas se cross_replica for usado.
  • num_partitions se cross_partition for usado.
• (C3) type(result) = type(operand).

Exemplos

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

collective_permute

Semântica

Em cada grupo de processos na grade de processos do StableHLO, envia o valor do tensor operand do processo de origem para o de destino e produz um tensor result.

A operação divide a grade de processo do StableHLO em process_groups, que é definida da seguinte forma:

• cross_replica(source_target_pairs) se channel_id <= 0.
• cross_partition(source_target_pairs) se channel_id > 0.

Depois disso, result@process é dado por:

• operand@process_groups[i, 0], se houver um i em que process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) caso contrário.

Entradas

Saídas

Restrições

• (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, em que N é definido como:
  • num_replicas se cross_replica for usado.
  • num_partitions se cross_partition for usado.
• (C5) type(result) = type(operand).

Exemplos

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

 Mais exemplos

compare

Semântica

Realiza uma comparação elemento a elemento dos tensores lhs e rhs de acordo com comparison_direction e compare_type e produz um tensor result.

Os valores de comparison_direction e compare_type têm a seguinte semântica:

Para tipos de elementos booleanos e inteiros:

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

Para tipos de elementos de ponto flutuante com compare_type = FLOAT, a operação implementa as seguintes operações IEEE-754:

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

Para tipos de elementos de ponto flutuante com compare_type = TOTALORDER, a operação usa a combinação de operações totalOrder e compareQuietEqual do IEEE-754.

Para tipos de elementos complexos, a comparação lexicográfica de pares (real, imag) é realizada usando o comparison_direction e o compare_type fornecidos. Impor uma ordenação em números complexos envolve semântica surpreendente. Por isso, no futuro, planejamos remover o suporte a números complexos quando comparison_direction for GE, GT, LE ou LT (#560).

Para tipos quantizados, realiza dequantize_compare(lhs, rhs, comparison_direction).

Entradas

Saídas

Restrições

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type é definido como:
  • SIGNED se is_signed_integer(element_type(lhs)).
  • UNSIGNED se is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT ou TOTALORDER se is_float(element_type(lhs)).
  • FLOAT se is_complex(element_type(lhs)).

Exemplos

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

 Mais exemplos

complexo

Semântica

Realiza a conversão de elemento por elemento para um valor complexo de um par de valores reais e imaginários, lhs e rhs, e produz um tensor result.

Entradas

Saídas

Restrições

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) tem o tipo complex<E> em que E = element_type(lhs).

Exemplos

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

 Mais exemplos

composto

Semântica

Encapsula uma operação composta por outras operações do StableHLO, usando inputs e composite_attributes e produzindo results. A semântica da operação é implementada pelo atributo decomposition. A operação composite pode ser substituída pela decomposição dela sem mudar a semântica do programa. Nos casos em que a inclusão inline da decomposição não fornece a mesma semântica de operação, prefira usar custom_call.

O campo version (padrão 0) é usado para indicar quando a semântica de um composto muda.

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

concatenate

Semântica

Concatena inputs ao longo da dimensão dimension na mesma ordem dos argumentos fornecidos e produz um tensor result. De maneira mais formal, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], em que:

1. id = d0 + ... + dk-1 + kd.
2. d é igual a dimension, e d0, ... são tamanhos da dª dimensão de inputs.

Entradas

Saídas

Restrições

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) exceto para 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]), exceto:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Exemplos

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

 Mais exemplos

constante

Semântica

Produz um tensor output de uma constante value.

Entradas

Saídas

Restrições

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

Exemplos

%output = "stablehlo.constant"() {
  val<ue = dense[[0.0, 1.0], [>2.0, 3.0]<] : ten>sor2x2xf3>2
} : (<) - ten>sor2x2xf32
// %output: [[0.0, 1.0], [2.0, 3.0]]

 Mais exemplos

fazer uma conversão

Semântica

Executa uma conversão elemento a elemento de um tipo de elemento para outro no tensor operand e produz um tensor result.

Para conversões boolean-to-any-supported-type, o valor false é convertido em zero, e o valor true é convertido em um. Para conversões de any-supported-type-to-boolean, um valor zero é convertido em false, e valores diferentes de zero são convertidos em true. Confira abaixo como isso funciona para tipos complexos.

Para conversões envolvendo inteiro para inteiro, inteiro para ponto flutuante ou ponto flutuante para ponto flutuante, se o valor de origem puder ser representado exatamente no tipo de destino, o valor resultante será essa representação exata. Caso contrário, o comportamento será definido (#180).

Para conversões que envolvem floating-point-to-integer, a parte fracionária é truncada. Se o valor truncado não puder ser representado no tipo de destino, o comportamento será definido posteriormente (#180).

A conversão envolvendo complexo para complexo segue o mesmo comportamento das conversões de ponto flutuante para ponto flutuante para converter partes reais e imaginárias.

Para conversões complex-to-any-other-type e any-other-type-to-complex, o valor imaginário de origem é ignorado ou o valor imaginário de destino é zerado, respectivamente. A conversão da parte real segue as conversões de ponto flutuante.

Em princípio, essa operação pode expressar a desquantização (conversão de tensores quantizados em tensores regulares), a quantização (conversão de tensores regulares em tensores quantizados) e a requantização (conversão entre tensores quantizados), mas no momento temos operações dedicadas para isso: uniform_dequantize para o primeiro caso de uso e uniform_quantize para o segundo e o terceiro. No futuro, essas duas operações poderão ser combinadas em convert (#1576).

Entradas

Saídas

Restrições

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

Exemplos

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand)< : (t>ens>or3xi64<) - tenso<r3x>>complexf64
// %result: [(-1.0, 0.0), (0.0, 0.0), (1.0, 0.0)]

 Mais exemplos

convolução

Semântica

Calcula produtos escalares entre janelas de lhs e intervalos de rhs e gera result. O diagrama a seguir mostra como os elementos em result são calculados com base em lhs e rhs usando um exemplo concreto.

De maneira mais formal, considere o seguinte reformulação das entradas em termos de lhs para expressar janelas de 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).

Essa reformulação usa as seguintes funções auxiliares:

• 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] em que j[d] = i[permutation[d]].

Se feature_group_count = 1 e batch_group_count = 1, então para todo output_spatial_index em index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product em que:

• 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]). Esse recurso parece não estar sendo usado. Por isso, planejamos removê-lo no futuro (#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]).

Se 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).

Se 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).

Para tipos quantizados, realiza 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)).

Para tipos quantizados híbridos, realiza 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).

Entradas

Saídas

Restrições

• (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) Dado 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) Dado 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) Dado 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) é definido como:
  • dim(lhs, input_batch_dimension) / batch_group_count se result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) se result_dim = output_feature_dimension.
  • num_windows caso contrário, em que:
  • 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.
• Se a operação usar tensores não quantizados:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Se a operação usar tensores quantizados:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Se is_per_axis_quantized(rhs), então quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Se is_per_axis_quantized(result), então quantization_dimension(result) = output_feature_dimension.
  • Se is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Se is_per_tensor_quantized(rhs), então is_per_tensor_quantized(result).
  • Se !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Exemplos

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

 Mais exemplos

cosseno

Semântica

Executa a operação de cosseno elemento a elemento no tensor operand e gera um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para pontos flutuantes: cos do IEEE-754.
• Para números complexos: cosseno complexo.
• Para tipos quantizados: dequantize_op_quantize(cosine, operand, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

count_leading_zeros

Semântica

Realiza a contagem de elementos do número de bits zero à esquerda no tensor operand e produz um tensor result.

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

custom_call

Semântica

Encapsula uma operação call_target_name definida pela implementação que usa inputs e called_computations e produz results. has_side_effect, backend_config e api_version podem ser usados para fornecer metadados adicionais definidos pela implementação.

No momento, essa operação contém uma coleção bastante desorganizada de metadados que reflete a evolução orgânica da operação equivalente no compilador XLA. No futuro, planejamos unificar esses metadados (#741).

Entradas

Saídas

(Suporte a GPU XLA) Destinos especiais de custom_call

Há três call_target_name especiais relacionados a tipos buffer: CreateBuffer cria um buffer não inicializado, Pin cria um buffer inicializado e Unpin desaloca um buffer e retorna o conteúdo do buffer.

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

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

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

Alias

Algumas operações custom_call podem exigir uma parte nas saídas e outra nos operandos para compartilhar a mesma memória. Isso pode ser expresso por output_operand_aliases. Uma representação de par de alias consiste em uma lista de índices de tuplas de saída que representam a parte de saída e um operand_index com uma lista de índices de tuplas de operando que representam a parte de operando. A lista de índices de tuplas de saída ou de operandos fica vazia se o tipo correspondente não for tuple e pode ser arbitrariamente longa para um tipo de tupla arbitrariamente aninhado. Isso é semelhante à representação do alias XLA.

As partes de saída e entrada em um par de alias precisam ter o mesmo tipo. Para operações custom_call que não são chamadas para CreateBuffer, Pin e Unpin, um operando buffer pode aparecer em no máximo um par de alias, e uma saída buffer precisa aparecer em um par de alias.

Exemplos

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

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

dividir

Semântica

Realiza a divisão elemento a elemento dos tensores dividendo lhs e divisor rhs e gera um tensor result. Dependendo do tipo de elemento, faça o seguinte:

• Para números inteiros: divisão inteira que produz o quociente algébrico com qualquer parte fracionária descartada.
• Para pontos flutuantes: division do IEEE-754.
• Para números complexos: divisão complexa.
• Para tipos quantizados:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Entradas

Saídas

Restrições

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

Exemplos

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

 Mais exemplos

dot_general

Semântica

Calcula produtos escalares entre fatias de lhs e fatias de rhs e produz um tensor result.

Mais formalmente, result[result_index] = dot_product, em que:

• 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 em que size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) e 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)).

Para tipos quantizados, realiza 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)).

Para tipos quantizados híbridos, realiza 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 controla a compensação entre velocidade e acurácia para computações em back-ends de aceleradores. Pode ser um dos seguintes valores (no momento, a semântica desses valores de enumeração está subespecificada, mas planejamos resolver isso em #755):

• DEFAULT: cálculo mais rápido, mas aproximação menos precisa do número original.
• HIGH: cálculo mais lento, mas aproximação mais precisa do número original.
• HIGHEST: cálculo mais lento, mas aproximação mais precisa do número original.

Um DotAlgorithm define as principais propriedades do algoritmo usado para implementar a operação de ponto, que também define a precisão. Se os campos de atributo do algoritmo forem definidos, o precision_config precisará ser DEFAULT. DotAlgorithms não têm um valor padrão, já que os parâmetros padrão são definidos pela implementação. Assim, todos os campos do algoritmo de ponto podem ser definidos como None para especificar um algoritmo de ponto vazio, que usará o valor precision_config.

Os campos DotAlgorithm incluem:

• lhs_precision_type e rhs_precision_type, as precisões para as quais os lados esquerdo e direito da operação são arredondados. Os tipos de precisão são independentes dos tipos de armazenamento das entradas e da saída.
• accumulation_type a precisão usada para acumulação.
• lhs_component_count, rhs_component_count e num_primitive_operations são aplicados quando usamos um algoritmo que decompõe o lado esquerdo e/ou direito em vários componentes e faz várias operações de ponto "primitivas" nesses valores, geralmente para emular uma precisão maior (por exemplo, Como aproveitar o tipo de dados de inteligência artificial bfloat16 para cálculos de maior precisão: bf16_6x tf32_3x etc.). Para algoritmos sem decomposição, esses valores precisam ser definidos como 1.
• allow_imprecise_accumulation para especificar se o acúmulo em precisão menor é permitido para algumas etapas (por exemplo, CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Exemplos de atributos DotAlgorithm:

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


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


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

Cabe às implementações decidir quais combinações são compatíveis. Em geral, não há garantia de que cada algoritmo seja compatível com cada tipo de acelerador pelo consumidor do StableHLO. Se um determinado algoritmo não for compatível, um erro deverá ser gerado em vez de usar uma alternativa. A verificação do StableHLO vai fornecer a melhor verificação possível, evitando algoritmos que não são conhecidos por serem compatíveis com qualquer hardware.

Consulte xla_data.proto > Algorithm para conferir alguns valores de algoritmo compatíveis. O tíquete nº 2483 registra o plano de criar um documento centralizado sobre algoritmos compatíveis por back-end.

Entradas

Saídas

Restrições

• (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).
• Se a operação usar tensores não quantizados:
  • (C13) element_type(lhs) = element_type(rhs).
• Se a operação usar tensores quantizados:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Se is_per_axis_quantized(rhs), então quantization_dimension(rhs) não está em rhs_contracting_dimensions.
  • Se is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Se is_per_tensor_quantized(rhs), então is_per_tensor_quantized(result).
  • Se !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Se !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation):
  • (C21) precision_config... = DEFAULT.
  • (C22) 0 < lhs_component_count.
  • (C23) 0 < rhs_component_count.
  • (C24) 0 < num_primitive_operations.

Exemplos

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

 Mais exemplos

dynamic_broadcast_in_dim

Semântica

Essa operação é funcionalmente idêntica à operação broadcast_in_dim, mas o formato do resultado é especificado dinamicamente via output_dimensions.

A operação também aceita atributos opcionais known_expanding_dimensions, known_nonexpanding_dimensions para expressar conhecimento estático sobre o comportamento de expansão das dimensões. Se não for especificado, todas as dimensões serão consideradas expansíveis.

Entradas

Saídas

Restrições

• (C1) element_type(result) é dado por:
  • element_type(operand), se !is_per_axis_quantized(operand).
  • element_type(operand), exceto que quantization_dimension(operand), scales(operand) e zero_points(operand) podem ser diferentes de quantization_dimension(result), scales(result) e zero_points(result), respectivamente, caso contrário.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Para todos os d em axes(operand):
  • dim(operand, d) = 1 ou
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Se is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Se dim(operand, quantization_dimension(operand)) = 1, então scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).
• (C7) size(output_dimensions) = rank(result).
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand).

Exemplos

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

 Mais exemplos

dynamic_conv

Semântica

Essa operação é funcionalmente idêntica à operação de convolução, mas o padding é especificado dinamicamente via padding.

Entradas