StableHLO-Spezifikation

StableHLO ist eine Reihe von Operationen für High-Level-Operationen (HLO) in Modellen für maschinelles Lernen (ML). StableHLO dient als Portabilitätsebene zwischen verschiedenen ML-Frameworks und ML-Compilern: ML-Frameworks, die StableHLO-Programme erstellen, sind mit ML-Compilern kompatibel, die StableHLO-Programme verwenden.

Unser Ziel ist es, die ML-Entwicklung zu vereinfachen und zu beschleunigen, indem wir die Interoperabilität zwischen verschiedenen ML-Frameworks (z. B. TensorFlow, JAX und PyTorch) und ML-Compilern (z. B. XLA und IREE) verbessern. Zu diesem Zweck enthält dieses Dokument eine Spezifikation für die StableHLO-Programmiersprache.

Diese Spezifikation enthält drei Hauptabschnitte. Zuerst wird im Abschnitt Programme die Struktur von StableHLO-Programmen beschrieben, die aus StableHLO-Funktionen bestehen, die wiederum aus StableHLO-Operationen bestehen. Im Abschnitt Ops dieser Struktur wird die Semantik der einzelnen Vorgänge angegeben. Im Abschnitt Ausführung finden Sie die Semantik für alle diese Vorgänge, die gemeinsam in einem Programm ausgeführt werden. Im Abschnitt Notation wird die in der gesamten Spezifikation verwendete Notation erläutert.

Wenn Sie die Spezifikation einer früheren Version von StableHLO aufrufen möchten, öffnen Sie das Repository für das entsprechende Release mit Tag. Ein Beispiel ist die StableHLO v0.19.0-Spezifikation. Änderungen, die bei jeder Nebenversionserhöhung von StableHLO aufgetreten sind, finden Sie im Versionslog in VhloDialect.td.

Programme

Program ::= {Func}

StableHLO-Programme bestehen aus einer beliebigen Anzahl von StableHLO-Funktionen. Unten sehen Sie ein Beispielprogramm mit einer Funktion @main, die drei Eingaben (%image, %weights und %bias) und eine Ausgabe hat. Der Funktionskörper hat sechs Vorgänge.

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

Funktionen

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

StableHLO-Funktionen (auch benannte Funktionen genannt) haben eine Kennung, Ein-/Ausgaben und einen Rumpf. In Zukunft planen wir, zusätzliche Metadaten für Funktionen einzuführen, um eine bessere Kompatibilität mit HLO zu erreichen (#425, #626, #740, #744).

IDs

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

StableHLO-Kennzeichnungen ähneln Kennzeichnungen in vielen Programmiersprachen, haben aber zwei Besonderheiten: 1) Alle Kennzeichnungen haben Sigils, die verschiedene Arten von Kennzeichnungen unterscheiden, 2) Wertkennzeichnungen können vollständig numerisch sein, um die Generierung von StableHLO-Programmen zu vereinfachen.

Typen

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

StableHLO-Typen werden in Werttypen (auch Erstklassentypen genannt), die StableHLO-Werte darstellen, und Nicht-Werttypen, die andere Programmelemente beschreiben, unterteilt. StableHLO-Typen ähneln Typen in vielen Programmiersprachen. Die Besonderheit ist, dass StableHLO domänenspezifisch ist, was zu einigen ungewöhnlichen Ergebnissen führt (z. B. sind skalare Typen keine Wertetypen).

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

Tensortypen stellen Tensoren dar, d.h. mehrdimensionale Arrays. Sie haben eine Form und einen Elementtyp. Eine Form stellt nicht negative oder unbekannte Dimensionsgrößen in aufsteigender Reihenfolge der entsprechenden Dimensionen (auch Achsen genannt) dar, die von 0 bis R-1 nummeriert sind. Die Anzahl der Dimensionen R wird als Rang bezeichnet. tensor<2x3xf32> ist beispielsweise ein Tensortyp mit der Form 2x3 und dem Elementtyp f32. Sie hat zwei Dimensionen (oder anders ausgedrückt: zwei Achsen) – die 0. und die 1. Dimension – mit den Größen 2 und 3. Der Rang ist 2.

Formen können teilweise oder vollständig unbekannt sein (dynamisch), z.B. ist tensor<?x2xf64> teilweise unbekannt und tensor<?x?xf64> vollständig unbekannt. Dynamische Dimensionsgrößen werden mit einem ? dargestellt. Formen können nicht aus dem Ranking entfernt werden.

In Zukunft möchten wir die Tensortypen über Dimensionsgrößen und Elementtypen hinaus erweitern, z. B. um Layouts (#629) und Sparsity (#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

Quantisierte Elementtypen stellen Ganzzahlwerte eines Speichertyp im Bereich von storage_min bis storage_max (einschließlich) dar, die Gleitkommawerten eines ausgedrückten Typs entsprechen. Für einen bestimmten Ganzzahlwert i kann der entsprechende Gleitkommawert f als f = (i - zero_point) * scale berechnet werden, wobei scale und zero_point als Quantisierungsparameter bezeichnet werden. storage_min und storage_max sind in der Grammatik optional, haben aber die Standardwerte min_value(storage_type) bzw. max_value(storage_type). Für quantisierte Elementtypen gelten die folgenden Einschränkungen:

• (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) Wenn is_empty(quantization_dimension), dann size(scales) = 1.
• (C11) 0 <= quantization_dimension.

Derzeit ist QuantizationScale eine Gleitkommakonstante. Es besteht jedoch großes Interesse an ganzzahligen Skalen, die mit Multiplikatoren und Verschiebungen dargestellt werden. Wir planen, dies in naher Zukunft zu untersuchen (#1404).

Es gibt eine laufende Diskussion über die Semantik von QuantizationZeroPoint, einschließlich des Typs, der Werte und der Frage, ob es nur einen oder möglicherweise mehrere Nullpunkte in einem quantisierten Tensortyp geben kann. Basierend auf den Ergebnissen dieser Diskussion kann sich die Spezifikation für null Punkte in Zukunft ändern (#1405).

Eine weitere laufende Diskussion betrifft die Semantik von QuantizationStorageMin und QuantizationStorageMax, um festzustellen, ob Einschränkungen für diese Werte und für die Werte quantisierter Tensoren auferlegt werden sollten (#1406).

Schließlich planen wir, unbekannte Skalen und Nullpunkte ähnlich darzustellen wie unbekannte Dimensionsgrößen (#1407).

Quantisierte Tensortypen stellen Tensoren mit quantisierten Elementen dar. Diese Tensoren sind genau wie reguläre Tensoren, nur dass ihre Elemente quantisierte Elementtypen anstelle von regulären Elementtypen haben.

Bei quantisierten Tensoren kann die Quantisierung pro Tensor erfolgen. Das bedeutet, dass für den gesamten Tensor ein scale- und ein zero_point-Wert vorhanden ist. Sie kann aber auch pro Achse erfolgen. Das bedeutet, dass mehrere scales- und zero_points-Werte vorhanden sind, ein Paar pro Slice einer bestimmten Dimension quantization_dimension. Formaler ausgedrückt: In einem Tensor t mit achsenweiser Quantisierung gibt es dim(t, quantization_dimension) Slices von quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] usw. Alle Elemente im i-ten Slice verwenden scales[i] und zero_points[i] als Quantisierungsparameter. Für quantisierte Tensortypen gelten die folgenden Einschränkungen:

• Für die Quantisierung pro Tensor:
  • Keine zusätzlichen Einschränkungen.
• Für die Quantisierung pro Achse:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Tokentypen stellen Tokens dar, d.h. undurchsichtige Werte, die von einigen Vorgängen erstellt und verwendet werden. Tokens werden verwendet, um die Ausführungsreihenfolge von Vorgängen festzulegen, wie im Abschnitt Ausführung beschrieben.

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

Puffertypen stellen Puffer dar. In XLA sind Puffer beispielsweise mehrdimensionale Arrays mit konsistentem Speicher. Ähnlich wie bei Tensortypen haben Puffertypen eine Form und einen Elementtyp. Eine Form stellt nicht negative oder unbekannte Dimensionsgrößen in aufsteigender Reihenfolge der entsprechenden Dimensionen (auch Achsen genannt) dar, die von 0 bis R-1 nummeriert sind. Die Anzahl der Dimensionen R wird als Rang bezeichnet. memref<2x3xf32> ist beispielsweise ein Puffertyp mit der Form 2x3 und dem Elementtyp f32. Es hat zwei Dimensionen (oder mit anderen Worten zwei Achsen) – die 0. Dimension und die 1. Dimension – mit den Größen 2 und 3. Der Rang ist 2.

Puffer können mit einem custom_call bis CreateBuffer oder Pin zugewiesen und mit einem custom_call bis Unpin freigegeben werden. Nur custom_call-Vorgänge können den Inhalt von Puffern lesen und schreiben. Weitere Informationen finden Sie unter custom_call.

Tupeltypen stellen Tupel dar, d.h. heterogene Listen. Tupel sind eine Legacy-Funktion, die nur zur Kompatibilität mit HLO vorhanden ist. In HLO werden Tupel verwendet, um variadische Ein- und Ausgaben darzustellen. In StableHLO werden variadische Ein- und Ausgaben nativ unterstützt. Tupel werden in StableHLO nur verwendet, um die HLO-ABI umfassend darzustellen, wobei z. B. T, tuple<T> und tuple<tuple<T>> je nach Implementierung erheblich variieren können. In Zukunft planen wir, Änderungen an der HLO-ABI vorzunehmen, die es uns ermöglichen, Tupeltypen aus StableHLO zu entfernen (#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'

Elementtypen stellen Elemente von Tensortypen dar. Anders als in vielen Programmiersprachen sind diese Typen in StableHLO nicht erstklassig. Das bedeutet, dass StableHLO-Programme Werte dieser Typen nicht direkt darstellen können. Daher ist es üblich, skalare Werte vom Typ T mit 0-dimensionalen Tensorwerten vom Typ tensor<T> darzustellen.

• Der boolesche Typ stellt die booleschen Werte true und false dar.
• Ganzzahltypen können entweder mit Vorzeichen (si) oder ohne Vorzeichen (ui) sein und eine der unterstützten Bitbreiten (2, 4, 8, 16, 32 oder 64) haben. Mit Vorzeichen versehene siN-Typen stellen Ganzzahlwerte von -2^(N-1) bis einschließlich 2^(N-1)-1 dar und vorzeichenlose uiN-Typen stellen Ganzzahlwerte von 0 bis einschließlich 2^N-1 dar.
• Gleitkommatypen können einen der folgenden Werte haben:
  • f8E3M4, f8E4M3 und f8E5M2 sind 8-Bit-Gleitkommazahlen gemäß IEEE-754-Konventionen.
  • Die Typen f8E4M3FN und f8E5M2 entsprechen den E4M3- und E5M2-Codierungen des FP8-Formats, die in FP8 Formats for Deep Learning beschrieben werden.
  • f8E4M3FNUZ- und f8E5M2FNUZ-Typen, die den E4M3- und E5M2-Codierungen der in 8-bit Numerical Formats for Deep Neural Networks beschriebenen FP8-Formate entsprechen.
  • f8E4M3B11FNUZ-Typ, der der E4M3-Codierung der FP8-Formate entspricht, die in Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks beschrieben werden.
  • bf16-Typ, der dem bfloat16-Format entspricht, das in BFloat16: The secret to high performance on Cloud TPUs beschrieben wird.
  • Die Typen f16, f32 und f64 entsprechen den Formaten binary16 („halbe Genauigkeit“), binary32 („einfache Genauigkeit“) und binary64 („doppelte Genauigkeit“), die im IEEE 754-Standard beschrieben werden.
  • Der Typ tf32 entspricht dem TensorFloat32-Format und wird in StableHLO nur eingeschränkt unterstützt.
  • f4E2M1FN-, f6E2M3FN-, f6E3M2FN- und f8E8M0FNU-MX-Typen (Mikroskalierung), die in der OCP Microscaling Formats Specification beschrieben werden.
• Komplexe Typen stellen komplexe Werte mit einem Realteil und einem Imaginärteil desselben Elementtyps dar. Unterstützte komplexe Typen sind complex<f32> (beide Teile sind vom Typ f32) und complex<f64> (beide Teile sind vom Typ f64).

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

Funktionstypen stellen sowohl benannte als auch anonyme Funktionen dar. Sie haben Eingabetypen (die Liste der Typen auf der linken Seite von ->) und Ausgabetypen (die Liste der Typen auf der rechten Seite von ->). In vielen Programmiersprachen sind Funktionstypen erstklassig, in StableHLO jedoch nicht.

StringType ::= 'string'

String-Typ steht für Bytefolgen. Im Gegensatz zu vielen Programmiersprachen ist der String-Typ in StableHLO nicht erstklassig und wird nur verwendet, um statische Metadaten für Programmelemente anzugeben.

Vorgänge

StableHLO-Vorgänge (auch Ops genannt) stellen eine geschlossene Menge von Vorgängen auf hoher Ebene in Machine-Learning-Modellen dar. Wie oben beschrieben, ist die StableHLO-Syntax stark von MLIR inspiriert. Das ist nicht unbedingt die ergonomischste Alternative, aber wohl die beste Lösung für das Ziel von StableHLO, die Interoperabilität zwischen ML-Frameworks und ML-Compilern zu verbessern.

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

StableHLO-Vorgänge (auch Ops genannt) haben einen Namen, Ein-/Ausgaben und eine Signatur. Der Name besteht aus dem Präfix stablehlo. und einem Mnemonik, das einen der unterstützten Vorgänge eindeutig identifiziert. Unten finden Sie eine vollständige Liste aller unterstützten Operationen.

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

Vorgänge nutzen Eingaben und erzeugen Ausgaben. Eingaben werden in Eingabewerte (während der Ausführung berechnet), Eingabefunktionen (statisch bereitgestellt, da Funktionen in StableHLO keine erstklassigen Werte sind) und Eingabeattribute (ebenfalls statisch bereitgestellt) unterteilt. Die Art der Ein- und Ausgaben, die von einem Vorgang verwendet und erzeugt werden, hängt von seinem Mnemonic ab. Der add-Vorgang verwendet beispielsweise 2 Eingabewerte und erzeugt 1 Ausgabewert. Im Vergleich dazu benötigt der Vorgang select_and_scatter drei Eingabewerte, zwei Eingabefunktionen und drei Eingabeattribute.

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

Eingabefunktionen (auch anonyme Funktionen genannt) ähneln benannten Funktionen sehr, mit folgenden Ausnahmen: 1) Sie haben keine Kennung (daher der Name „anonym“), 2) Sie deklarieren keine Ausgabetypen (Ausgabetypen werden vom return-Vorgang innerhalb der Funktion abgeleitet).

Die Syntax für Eingabefunktionen enthält einen derzeit nicht verwendeten Teil (siehe die Unused-Produktion oben), der aus Gründen der Kompatibilität mit MLIR vorhanden ist. In MLIR gibt es ein allgemeineres Konzept von „Regionen“, die mehrere „Blöcke“ von Vorgängen haben können, die über Sprungvorgänge miteinander verbunden sind. Diese Blöcke haben IDs, die der Unused-Produktion entsprechen, sodass sie voneinander unterschieden werden können. StableHLO hat keine Jump-Operationen, daher wird der entsprechende Teil der MLIR-Syntax nicht verwendet (ist aber weiterhin vorhanden).

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

Eingabeattribute haben einen Namen und einen Wert, der eine der unterstützten Konstanten ist. Sie sind die primäre Methode zum Angeben statischer Metadaten für Programmelemente. Der concatenate-Vorgang verwendet beispielsweise das Attribut dimension, um die Dimension anzugeben, entlang derer die Eingabewerte verkettet werden. Ebenso verwendet der slice-Vorgang mehrere Attribute wie start_indices und limit_indices, um die Grenzen anzugeben, die zum Aufteilen des Eingabewerts verwendet werden.

Derzeit enthalten StableHLO-Programme manchmal Attribute, die in diesem Dokument nicht beschrieben werden. In Zukunft planen wir, diese Attribute entweder in das StableHLO-Opset aufzunehmen oder zu verhindern, dass sie in StableHLO-Programmen vorkommen. In der Zwischenzeit finden Sie hier die Liste dieser Attribute:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Standortmetadaten (#594).

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

Die Op-Signatur besteht aus den Typen aller Eingabewerte (der Liste der Typen auf der linken Seite von ->) und den Typen aller Ausgabewerte (der Liste der Typen auf der rechten Seite von ->). Streng genommen sind Eingabetypen redundant und Ausgabetypen sind fast immer redundant, da für die meisten StableHLO-Vorgänge Ausgabetypen aus Eingaben abgeleitet werden können. Die Signatur des Vorgangs ist jedoch bewusst Teil der StableHLO-Syntax, um die Kompatibilität mit MLIR zu gewährleisten.

Unten sehen Sie ein Beispiel für einen Vorgang mit dem Mnemonic select_and_scatter. Sie verwendet 3 Eingabewerte (%operand, %source und %init_value), 2 Eingabefunktionen und 3 Eingabeattribute (window_dimensions, window_strides und padding). Beachten Sie, dass die Signatur des Vorgangs nur die Typen der Eingabewerte enthält (nicht die Typen der Eingabefunktionen und ‑attribute, die inline angegeben werden).

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

Konstanten

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

StableHLO-Konstanten haben ein Literal und einen Typ, die zusammen einen StableHLO-Wert darstellen. Im Allgemeinen ist der Typ Teil der Konstantensyntax, es sei denn, er ist eindeutig (z.B. hat eine boolesche Konstante eindeutig den Typ i1, während eine Ganzzahlkonstante mehrere mögliche Typen haben kann).

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

Boolesche Konstanten stellen die booleschen Werte true und false dar. Boolesche Konstanten haben den Typ i1.

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

Ganzzahlkonstanten stellen Ganzzahlwerte über Strings dar, die die Dezimal- oder Hexadezimalnotation verwenden. Andere Basen, z.B. binär oder oktal, werden nicht unterstützt. Für Ganzzahlkonstanten gelten die folgenden Einschränkungen:

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

Gleitkommakonstanten stellen Gleitkommawerte über Strings dar, die die Dezimal- oder wissenschaftliche Notation verwenden. Außerdem kann die hexadezimale Notation verwendet werden, um die zugrunde liegenden Bits im Gleitkommaformat des entsprechenden Typs direkt anzugeben. Für Gleitkommakonstanten gelten die folgenden Einschränkungen:

• (C1) Wenn keine hexadezimale Notation verwendet wird,is_wellformed(float_literal, float_type).
• (C2) Wenn die hexadezimale Notation verwendet wird, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Komplexe Konstanten stellen komplexe Werte mithilfe von Listen mit einem Realteil (kommt zuerst) und einem Imaginärteil (kommt als Zweites) dar. Beispiel: (1.0, 0.0) : complex<f32> steht für 1.0 + 0.0i und (0.0, 1.0) : complex<f32> für 0.0 + 1.0i. Die Reihenfolge, in der diese Teile dann im Arbeitsspeicher gespeichert werden, ist implementierungsabhängig. Für komplexe Konstanten gelten die folgenden Einschränkungen:

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

Tensorkonstanten stellen Tensorwerte mithilfe von verschachtelten Listen dar, die über die NumPy-Notation angegeben werden. Beispielsweise stellt dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> einen Tensorwert mit der folgenden Zuordnung von Index zu Element dar: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. Die Reihenfolge, in der diese Elemente dann im Speicher abgelegt werden, ist implementierungsabhängig. Für Tensor-Konstanten gelten die folgenden Einschränkungen:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), wobei:
  • 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)), wobei:
  • 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:]).
  • Andernfalls false.

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

Quantisierte Tensorkonstanten stellen quantisierte Tensorwerte mit derselben Notation wie Tensorkonstanten dar. Die Elemente werden als Konstanten ihres Speichertyps angegeben. Für quantisierte Tensor-Konstanten gelten die folgenden Einschränkungen:

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

Stringliterale bestehen aus Byte, die mit ASCII-Zeichen und Escape-Sequenzen angegeben werden. Sie sind unabhängig von der Codierung, daher ist die Interpretation dieser Bytes implementierungsdefiniert. Stringliterale haben den Typ string.

Operativer Betrieb

abs

Semantik

Führt eine elementweise „abs“-Operation für den Tensor operand aus und erzeugt einen result-Tensor. Je nach Elementtyp wird Folgendes ausgeführt:

• Bei Ganzzahlen mit Vorzeichen: Ganzzahlmodul.
• Für Gleitkommazahlen: abs aus IEEE-754.
• Für komplexe Zahlen: komplexer Modulus.
• Für quantisierte Typen: dequantize_op_quantize(abs, operand, type(result)).

Eingaben

Ausgaben

Einschränkungen

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) wird so definiert:
  • complex_element_type(element_type(operand)), wenn is_complex(operand).
  • Andernfalls baseline_element_type(operand).

Beispiele

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

 Weitere Beispiele

Hinzufügen

Semantik

Führt eine elementweise Addition von zwei Tensoren lhs und rhs aus und erzeugt einen result-Tensor. Je nach Elementtyp wird Folgendes ausgeführt:

• Für boolesche Werte: logisches ODER.
• Bei Ganzzahlen: Addition von Ganzzahlen.
• Für Gleitkommazahlen: addition aus IEEE-754.
• Für komplexe Zahlen: komplexe Addition.
• Für quantisierte Typen: dequantize_op_quantize(add, lhs, rhs, type(result)).

Eingaben

Ausgaben

Einschränkungen

• Wenn für die Operation nicht quantisierte Tensoren verwendet werden:
  • (C1) type(lhs) = type(rhs) = type(result).
• Wenn für den Vorgang quantisierte Tensoren verwendet werden:
  • (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) Wenn is_per_axis_quantized(lhs), dann quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Wenn is_per_axis_quantized(rhs), dann quantization_dimension(rhs) = quantization_dimension(result).

Beispiele

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

 Weitere Beispiele

after_all

Semantik

Sorgt dafür, dass die Vorgänge, die inputs erzeugen, vor allen Vorgängen ausgeführt werden, die von result abhängen. Die Ausführung dieses Vorgangs hat keine Auswirkungen. Er ist nur dazu da, Datenabhängigkeiten von result zu inputs herzustellen.

Eingaben

Ausgaben

Beispiele

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

 Weitere Beispiele

all_gather

Semantik

Verkettet in jeder Prozessgruppe im StableHLO-Prozessraster die Werte der operands-Tensoren aus jedem Prozess entlang all_gather_dim und erzeugt results-Tensoren.

Bei der Operation wird das StableHLO-Prozessraster in process_groups aufgeteilt, das so definiert ist:

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

Gehen Sie anschließend in jedem process_group so vor:

• operands...@receiver = [operand@sender for sender in process_group] für alle receiver in process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) für alle process in process_group.

Eingaben

Ausgaben

Einschränkungen

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) wird so definiert:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_replicas, wenn cross_replica_and_partition verwendet wird.
  • num_processes, wenn flattened_ids verwendet wird.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Wenn use_global_device_ids = true, dann channel_id > 0.
• (C6) type(results...) = type(operands...) außer:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Beispiele

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

 Weitere Beispiele

all_reduce

Semantik

Wendet in jeder Prozessgruppe im StableHLO-Prozessraster eine Reduzierungsfunktion computation auf die Werte der operands-Tensoren aus jedem Prozess an und gibt results-Tensoren aus.

Bei der Operation wird das StableHLO-Prozessraster in process_groups aufgeteilt, das so definiert ist:

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

Gehen Sie anschließend in jedem process_group so vor:

• results...@process[result_index] = exec(schedule) für einen binären Baum schedule, wobei:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule ist ein implementierungsdefinierter binärer Baum, dessen Inorder-Durchlauf to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])) ist.

Eingaben

Ausgaben

Einschränkungen

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) wird so definiert:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_replicas, wenn cross_replica_and_partition verwendet wird.
  • num_processes, wenn flattened_ids verwendet wird.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Wenn use_global_device_ids = true, dann channel_id > 0.
• (C5) computation hat den Typ (tensor<E>, tensor<E>) -> (tensor<E>), wobei is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Beispiele

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

 Weitere Beispiele

all_to_all

Semantik

In jeder Prozessgruppe im StableHLO-Prozessraster werden die Werte der operands-Tensoren entlang split_dimension in Teile aufgeteilt, die auf die Prozesse verteilt, entlang concat_dimension verkettet und results-Tensoren erzeugt. Bei der Operation wird das StableHLO-Prozessraster in process_groups aufgeteilt, das so definiert ist:

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

Gehen Sie anschließend in jedem process_group so vor:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) für alle sender in process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group], wobei receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Eingaben

Ausgaben

Einschränkungen

• (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) wird so definiert:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_partitions, wenn cross_partition verwendet wird.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...), außer wenn split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Beispiele

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

 Weitere Beispiele

und

Semantik

Führt eine elementweise AND-Operation für zwei Tensoren lhs und rhs aus und gibt einen result-Tensor zurück. Je nach Elementtyp wird Folgendes ausgeführt:

• Für boolesche Werte: logisches UND.
• Für Ganzzahlen: bitweises UND.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

atan2

Semantik

Führt eine elementweise atan2-Operation für die Tensoren lhs und rhs aus und gibt einen result-Tensor zurück. Je nach Elementtyp wird Folgendes ausgeführt:

• Für Gleitkommazahlen: atan2 aus IEEE-754.
• Für komplexe Zahlen: komplexer atan2.
• Für quantisierte Typen: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

batch_norm_grad

Semantik

Berechnet die Gradienten mehrerer Eingaben von batch_norm_training durch Backpropagation von grad_output und erzeugt die Tensoren grad_operand, grad_scale und grad_offset. Formaler kann dieser Vorgang als Zerlegung in vorhandene StableHLO-Vorgänge mit der folgenden Python-Syntax ausgedrückt werden:

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

Für quantisierte Typen wird 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)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale und grad_offset haben dieselbe baseline_element_type.
• (C3) operand, grad_output und grad_operand haben dieselbe Form.
• (C4) scale, mean, variance, grad_scale und grad_offset haben dieselbe Form.
• (C5) size(scale) = dim(operand, feature_index).

Beispiele

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

Semantik

Normalisiert den operand-Tensor über alle Dimensionen hinweg, mit Ausnahme der feature_index-Dimension, und erzeugt einen result-Tensor. Formaler kann dieser Vorgang als Zerlegung in vorhandene StableHLO-Vorgänge mit der folgenden Python-Syntax ausgedrückt werden:

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)

Für quantisierte Typen wird 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)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance und result haben dieselbe 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).

Beispiele

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

Semantik

Berechnet den Mittelwert und die Varianz über alle Dimensionen hinweg, mit Ausnahme der feature_index-Dimension, und normalisiert den operand-Tensor, wodurch die Tensoren output, batch_mean und batch_var entstehen. Formaler kann dieser Vorgang als Zerlegung in vorhandene StableHLO-Vorgänge mit der folgenden Python-Syntax ausgedrückt werden:

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

Für quantisierte Typen wird 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)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var und output haben dieselbe 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).

Beispiele

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

Semantik

Führt einen Bitcast-Vorgang für den operand-Tensor aus und erzeugt einen result-Tensor, in dem die Bits des gesamten operand-Tensors mit dem Typ des result-Tensors neu interpretiert werden.

Formal ausgedrückt: Bei E = element_type(operand), E' = element_type(result) und R = rank(operand) gilt Folgendes:

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

bits gibt die In-Memory-Darstellung eines bestimmten Werts zurück. Das Verhalten ist implementierungsdefiniert, da die genaue Darstellung von Tensoren und Elementtypen implementierungsdefiniert ist.

Eingaben

Ausgaben

Einschränkungen

• (C1) Angesichts von E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) und R = rank(operand):
  • Wenn num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Wenn num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) für alle 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Wenn num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) für alle 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Wenn is_complex(operand) or is_complex(result), dann is_complex(operand) and is_complex(result).

Beispiele

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

 Weitere Beispiele

broadcast_in_dim

Semantik

Erweitert die Dimensionen und/oder den Rang eines Eingabetensors, indem die Daten im operand-Tensor dupliziert werden, und erzeugt einen result-Tensor. Formaler ausgedrückt: result[result_index] = operand[operand_index], wobei für alle d in axes(operand) gilt:

• operand_index[d] = 0, wenn dim(operand, d) = 1.
• Andernfalls operand_index[d] = result_index[broadcast_dimensions[d]].

Eingaben

Ausgaben

Einschränkungen

• (C1) element_type(result) wird so berechnet:
  • element_type(operand), wenn !is_per_axis_quantized(operand).
  • element_type(operand), mit der Ausnahme, dass quantization_dimension(operand), scales(operand) und zero_points(operand) von quantization_dimension(result), scales(result) und zero_points(result) abweichen können.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Für alle d in axes(operand):
  • dim(operand, d) = 1 oder
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Wenn is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Bei dim(operand, quantization_dimension(operand)) = 1, dann scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Beispiele

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

 Weitere Beispiele

Supportanfrage

Semantik

Gibt die Ausgabe zurück, die durch die Ausführung genau einer Funktion aus branches entsteht, je nach dem Wert von index. Formaler ausgedrückt: result = selected_branch(), wobei gilt:

• selected_branch = branches[index], wenn 0 <= index < size(branches).
• Andernfalls selected_branch = branches[-1].

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

cbrt

Semantik

Führt eine elementweise Kubikwurzeloperation für den operand-Tensor aus und erzeugt einen result-Tensor. Je nach Elementtyp wird Folgendes ausgeführt:

• Für Gleitkommazahlen: rootn(x, 3) aus IEEE-754.
• Für komplexe Zahlen: komplexe Kubikwurzel.
• Für quantisierte Typen: dequantize_op_quantize(cbrt, operand, type(result))

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

ceil

Semantik

Führt die elementweise Aufrundung des operand-Tensors aus und gibt einen result-Tensor aus. Implementiert den roundToIntegralTowardPositive-Vorgang aus der IEEE-754-Spezifikation. Für quantisierte Typen wird dequantize_op_quantize(ceil, operand, type(result)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

cholesky

Semantik

Berechnet die Cholesky-Zerlegung einer Reihe von Matrizen.

Formaler ausgedrückt: Für alle i in index_space(result) ist result[i0, ..., iR-3, :, :] eine Cholesky-Zerlegung von a[i0, ..., iR-3, :, :] in Form einer unteren Dreiecksmatrix (wenn lower gleich true ist) oder einer oberen Dreiecksmatrix (wenn lower gleich false ist). Die Ausgabewerte im gegenüberliegenden Dreieck, d.h. im strikten oberen oder strikten unteren Dreieck, sind implementierungsabhängig.

Wenn es i gibt, für das die Eingabematrix keine hermitesche positiv definite Matrix ist, ist das Verhalten nicht definiert.

Für quantisierte Typen wird dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

einschränken

Semantik

Beschneidet jedes Element des operand-Tensors auf einen Minimal- und Maximalwert und gibt einen result-Tensor aus. Formaler ausgedrückt: result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), wobei min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Für quantisierte Typen wird dequantize_op_quantize(clamp, min, operand, max, type(result)) ausgeführt.

Die Festlegung einer Reihenfolge für komplexe Zahlen birgt überraschende Semantiken. Daher planen wir, die Unterstützung für komplexe Zahlen für diesen Vorgang in Zukunft zu entfernen (#560).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

collective_broadcast

Semantik

Senden Sie in jeder Prozessgruppe im StableHLO-Prozessraster den Wert des operand-Tensors vom Quellprozess an die Zielprozesse und erstellen Sie einen result-Tensor.

Bei der Operation wird das StableHLO-Prozessraster in process_groups aufgeteilt, das so definiert ist:

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

Danach wird result@process so berechnet:

• operand@process_groups[i, 0], wenn es ein i gibt, sodass der Prozess in process_groups[i] ist.
• Andernfalls broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)).

Eingaben

Ausgaben

Einschränkungen

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N, wobei N so definiert ist:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_partitions, wenn cross_partition verwendet wird.
• (C3) type(result) = type(operand).

Beispiele

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

Semantik

In jeder Prozessgruppe im StableHLO-Prozessraster wird der Wert des Tensors operand vom Quellprozess an den Zielprozess gesendet und ein Tensor result erstellt.

Bei der Operation wird das StableHLO-Prozessraster in process_groups aufgeteilt, das so definiert ist:

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

Danach wird result@process so berechnet:

• operand@process_groups[i, 0], wenn ein i vorhanden ist, sodass process_groups[i, 1] = process.
• Andernfalls broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)).

Eingaben

Ausgaben

Einschränkungen

• (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, wobei N so definiert ist:
  • num_replicas, wenn cross_replica verwendet wird.
  • num_partitions, wenn cross_partition verwendet wird.
• (C5) type(result) = type(operand).

Beispiele

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

 Weitere Beispiele

vergleichen

Semantik

Führt einen elementweisen Vergleich der Tensoren lhs und rhs gemäß comparison_direction und compare_type durch und gibt einen result-Tensor aus.

Die Werte von comparison_direction und compare_type haben die folgende Semantik:

Für boolesche und Ganzzahlelementtypen:

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

Für Gleitkomma-Elementtypen mit compare_type = FLOAT implementiert der Vorgang die folgenden IEEE-754-Operationen:

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

Für Gleitkomma-Elementtypen mit compare_type = TOTALORDER verwendet der Vorgang die Kombination aus totalOrder- und compareQuietEqual-Vorgängen aus IEEE-754.

Bei komplexen Elementtypen wird ein lexikografischer Vergleich von (real, imag)-Paaren mit den bereitgestellten comparison_direction und compare_type durchgeführt. Die Festlegung einer Reihenfolge für komplexe Zahlen birgt überraschende Semantiken. Daher planen wir, die Unterstützung für komplexe Zahlen in Zukunft zu entfernen, wenn comparison_direction GE, GT, LE oder LT ist (#560).

Für quantisierte Typen wird dequantize_compare(lhs, rhs, comparison_direction) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type wird so definiert:
  • SIGNED, wenn is_signed_integer(element_type(lhs)).
  • UNSIGNED, wenn is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT oder TOTALORDER, wenn is_float(element_type(lhs)).
  • FLOAT, wenn is_complex(element_type(lhs)).

Beispiele

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

 Weitere Beispiele

komplex

Semantik

Führt die elementweise Konvertierung in einen komplexen Wert aus einem Paar von reellen und imaginären Werten, lhs und rhs, durch und erzeugt einen result-Tensor.

Eingaben

Ausgaben

Einschränkungen

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) hat den Typ complex<E>, wobei E = element_type(lhs).

Beispiele

// %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)]

 Weitere Beispiele

Zusammengesetzt

Semantik

Kapselt einen Vorgang, der aus anderen StableHLO-Vorgängen besteht, wobei inputs und composite_attributes verwendet und results ausgegeben werden. Die Semantik des Vorgangs wird durch das Attribut decomposition implementiert. Der Vorgang composite kann durch seine Zerlegung ersetzt werden, ohne dass sich die Programmsemantik ändert. Wenn durch das Inlining der Dekomposition nicht dieselbe Semantik für den Vorgang erreicht wird, sollten Sie custom_call verwenden.

Das Feld version (Standardwert: 0) wird verwendet, um anzugeben, wann sich die Semantik eines Composites ändert.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

%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

 Weitere Beispiele

concatenate

Semantik

Verkettet inputs entlang der Dimension dimension in derselben Reihenfolge wie die angegebenen Argumente und gibt einen result-Tensor aus. Formaler ausgedrückt: result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], wobei gilt:

1. id = d0 + ... + dk-1 + kd.
2. d entspricht dimension und d0, ... sind die Größen der d-ten Dimension von inputs.

Eingaben

Ausgaben

Einschränkungen

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) außer 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]) mit Ausnahme von:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Beispiele

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

 Weitere Beispiele

Konstante

Semantik

Erstellt einen output-Tensor aus einer konstanten value.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

eine Conversion ausführen

Semantik

Führt eine elementweise Konvertierung von einem Elementtyp in einen anderen für den operand-Tensor aus und gibt einen result-Tensor zurück.

Bei boolean-to-any-supported-type-Konvertierungen wird der Wert false in null und der Wert true in eins konvertiert. Bei any-supported-type-to-boolean-Konvertierungen wird ein Nullwert in false und ein Wert ungleich null in true konvertiert. Unten erfahren Sie, wie das bei komplexen Typen funktioniert.

Bei Konvertierungen, die Ganzzahl-zu-Ganzzahl, Ganzzahl-zu-Gleitkomma oder Gleitkomma-zu-Gleitkomma umfassen, ist der Ergebniswert die genaue Darstellung des Quellwerts im Zieltyp, sofern der Quellwert im Zieltyp genau dargestellt werden kann. Andernfalls ist das Verhalten noch nicht definiert (#180).

Bei Konvertierungen, bei denen floating-point-to-integer umgewandelt werden, wird der Bruchteil abgeschnitten. Wenn der gekürzte Wert nicht im Zieltyp dargestellt werden kann, ist das Verhalten noch nicht definiert (#180).

Bei der Konvertierung von komplex zu komplex wird das gleiche Verhalten wie bei der Konvertierung von Gleitkommazahl zu Gleitkommazahl verwendet, um Real- und Imaginärteile zu konvertieren.

Bei complex-to-any-other-type- und any-other-type-to-complex-Konvertierungen wird der imaginäre Quellwert ignoriert bzw. der imaginäre Zielwert auf null gesetzt. Die Konvertierung des Realteils folgt den Gleitkommakonvertierungen.

Grundsätzlich könnte dieser Vorgang die Dequantisierung (Konvertierung von quantisierten Tensoren in reguläre Tensoren), die Quantisierung (Konvertierung von regulären Tensoren in quantisierte Tensoren) und die Requantisierung (Konvertierung zwischen quantisierten Tensoren) ausdrücken. Derzeit haben wir jedoch dedizierte Vorgänge dafür: uniform_dequantize für den ersten Anwendungsfall und uniform_quantize für den zweiten und dritten Anwendungsfall. In Zukunft werden diese beiden Vorgänge möglicherweise in convert zusammengeführt (#1576).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

// %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)]

 Weitere Beispiele

Faltung

Semantik

Berechnet Punktprodukte zwischen Fenstern von lhs und Slices von rhs und gibt result aus. Das folgende Diagramm zeigt anhand eines konkreten Beispiels, wie Elemente in result aus lhs und rhs berechnet werden.

Formaler ausgedrückt: Die Eingaben werden in Bezug auf lhs neu formuliert, um Zeiträume von lhs ausdrücken zu können:

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

Für diese Umformulierung werden die folgenden Hilfsfunktionen verwendet:

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

Wenn feature_group_count = 1 und batch_group_count = 1, dann für alle output_spatial_index in index_space(dim(result, output_spatial_dimensions...)): result[result_shape(:, output_spatial_index, :)] = dot_product, wobei:

• 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]). Diese Funktion wird offenbar nicht verwendet. Wir planen daher, sie in Zukunft zu entfernen (#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]).

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

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

Führt für quantisierte Typen 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)) aus.

Für hybride quantisierte Typen wird 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) ausgeführt.

Eingaben

Ausgaben

Einschränkungen

• (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) Angenommen, 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) Angesichts von 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) Angesichts von 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) wird so definiert:
  • dim(lhs, input_batch_dimension) / batch_group_count, wenn result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension), wenn result_dim = output_feature_dimension.
  • Andernfalls num_windows, wobei gilt:
  • 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.
• Wenn für die Operation nicht quantisierte Tensoren verwendet werden:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Wenn für den Vorgang quantisierte Tensoren verwendet werden:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Wenn is_per_axis_quantized(rhs), dann quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Wenn is_per_axis_quantized(result), dann quantization_dimension(result) = output_feature_dimension.
  • Wenn is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Wenn is_per_tensor_quantized(rhs), dann is_per_tensor_quantized(result).
  • Wenn !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Beispiele

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

 Weitere Beispiele

Kosinus

Semantik

Führt eine elementweise Kosinusoperation für den operand-Tensor aus und erzeugt einen result-Tensor. Je nach Elementtyp wird Folgendes ausgeführt:

• Für Gleitkommazahlen: cos aus IEEE-754.
• Für komplexe Zahlen: komplexer Kosinus.
• Für quantisierte Typen: dequantize_op_quantize(cosine, operand, type(result)).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

count_leading_zeros

Semantik

Führt eine elementweise Zählung der Anzahl der führenden Nullbits im operand-Tensor durch und erzeugt einen result-Tensor.

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

custom_call

Semantik

Kapselt einen implementierungsdefinierten Vorgang call_target_name, der inputs und called_computations verwendet und results erzeugt. has_side_effect, backend_config und api_version können verwendet werden, um zusätzliche implementierungsdefinierte Metadaten bereitzustellen.

Derzeit enthält dieser Vorgang eine ziemlich unorganisierte Sammlung von Metadaten, die die organische Entwicklung des entsprechenden Vorgangs im XLA-Compiler widerspiegelt. In Zukunft planen wir, diese Metadaten zu vereinheitlichen (#741).

Eingaben

Ausgaben

(XLA-GPU-Unterstützung) Spezielle custom_call-Ziele

Es gibt drei spezielle call_target_name im Zusammenhang mit buffer-Typen: CreateBuffer erstellt eine nicht initialisierte buffer, Pin erstellt eine initialisierte buffer und Unpin gibt eine buffer frei und gibt den Inhalt der buffer zurück.

%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

Bei einigen custom_call-Vorgängen muss ein Teil der Ausgaben und ein Teil der Operanden denselben Speicher verwenden. Das lässt sich mit output_operand_aliases ausdrücken. Eine Alias-Paar-Darstellung besteht aus einer Liste von Ausgabetupel-Indexen, die den Ausgabeteil darstellen, und einem operand_index zusammen mit einer Liste von Operanden-Tupel-Indexen, die den Operandenteil darstellen. Die Liste der Ausgabeargument- oder Operanden-Tupelindizes ist leer, wenn der entsprechende Typ kein tuple-Typ ist. Sie kann für einen beliebig geschachtelten Tupeltyp beliebig lang sein. Dies ähnelt der XLA-Aliasdarstellung.

Der Ausgabeteil und der Eingabeteil in einem Aliaspaar müssen denselben Typ haben. Bei custom_call-Vorgängen, die nicht CreateBuffer, Pin und Unpin aufrufen, kann ein buffer-Operand in höchstens einem Aliaspaar und eine buffer-Ausgabe in einem Aliaspaar vorkommen.

Beispiele

%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

Dividieren

Semantik

Führt die elementweise Division der Tensoren „dividend“ lhs und „divisor“ rhs aus und gibt einen result-Tensor zurück. Je nach Elementtyp wird Folgendes ausgeführt:

• Bei Ganzzahlen: Ganzzahldivision, bei der der algebraische Quotient ohne Bruchteil zurückgegeben wird.
• Für Gleitkommazahlen: division aus IEEE-754.
• Für komplexe Zahlen: komplexe Division.
• Für quantisierte Typen:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Eingaben

Ausgaben

Einschränkungen

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

Beispiele

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

 Weitere Beispiele

dot_general

Semantik

Berechnet die Skalarprodukte zwischen Slices von lhs und Slices von rhs und gibt einen result-Tensor aus.

Formaler ausgedrückt: result[result_index] = dot_product, wobei gilt:

• 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 wobei size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) und 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)).

Führt für quantisierte Typen 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)) aus.

Für hybride quantisierte Typen wird 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) ausgeführt.

precision_config steuert den Kompromiss zwischen Geschwindigkeit und Genauigkeit für Berechnungen auf Accelerator-Back-Ends. Dies kann einer der folgenden Werte sein (derzeit sind die Semantiken dieser Enum-Werte nicht genau festgelegt, aber wir planen, dies in #755 zu ändern):

• DEFAULT: Schnellste Berechnung, aber ungenaueste Annäherung an die ursprüngliche Zahl.
• HIGH: Langsamere Berechnung, aber genauere Annäherung an die ursprüngliche Zahl.
• HIGHEST: Langsamste Berechnung, aber genaueste Annäherung an die Originalzahl.

Ein DotAlgorithm definiert die Haupteigenschaften des Algorithmus, der zur Implementierung der Punktoperation verwendet wird. Dadurch wird auch die Genauigkeit definiert. Wenn die Felder für das Algorithmusattribut festgelegt sind, muss precision_config DEFAULT sein. DotAlgorithms haben keinen Standardwert, da die Standardparameter implementierungsdefiniert sind. Daher können alle Felder des Punktalgorithmus auf None gesetzt werden, um einen leeren Punktalgorithmus anzugeben, der stattdessen den Wert precision_config verwendet.

DotAlgorithm-Felder:

• lhs_precision_type und rhs_precision_type: Die Genauigkeiten, auf die die linke und rechte Seite des Vorgangs gerundet werden. Die Genauigkeitstypen sind unabhängig von den Speichertypen der Ein- und Ausgabe.
• accumulation_type die für die Akkumulierung verwendete Genauigkeit.
• lhs_component_count, rhs_component_count und num_primitive_operations werden angewendet, wenn wir einen Algorithmus verwenden, der die linke und/oder rechte Seite in mehrere Komponenten zerlegt und mehrere „primitive“ Punktoperationen für diese Werte ausführt, in der Regel um eine höhere Genauigkeit zu emulieren (z. B. bfloat16-Datentyp für künstliche Intelligenz für Berechnungen mit höherer Genauigkeit nutzen: bf16_6x tf32_3x usw.). Bei Algorithmen ohne Dekomposition sollten diese Werte auf 1 gesetzt werden.
• allow_imprecise_accumulation, um anzugeben, ob die Akkumulierung in niedrigerer Genauigkeit für einige Schritte zulässig ist (z.B. CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Beispielattribute für 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}

Es liegt an den Implementierungen, welche Kombinationen unterstützt werden. Im Allgemeinen kann nicht garantiert werden, dass jeder Algorithmus von jedem Beschleunigertyp unterstützt wird. Wenn ein bestimmter Algorithmus nicht unterstützt wird, sollte ein Fehler ausgegeben werden, anstatt auf eine Alternative zurückzugreifen. Die StableHLO-Überprüfung erfolgt nach dem Best-Effort-Prinzip und verhindert die Verwendung von Algorithmen, die auf keiner Hardware unterstützt werden.

Einige unterstützte Algorithmuswerte finden Sie unter xla_data.proto > Algorithm. Im Ticket 2483 wird der Plan beschrieben, ein zentrales Dokument zu den unterstützten Algorithmen nach Backend zu erstellen.

Eingaben

Ausgaben

Einschränkungen

• (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).
• Wenn für die Operation nicht quantisierte Tensoren verwendet werden:
  • (C13) element_type(lhs) = element_type(rhs).
• Wenn für den Vorgang quantisierte Tensoren verwendet werden:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Wenn is_per_axis_quantized(rhs), dann quantization_dimension(rhs) nicht in rhs_contracting_dimensions.
  • Wenn is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Wenn is_per_tensor_quantized(rhs), dann is_per_tensor_quantized(result).
  • Wenn !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Wenn !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.

Beispiele

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

 Weitere Beispiele

dynamic_broadcast_in_dim

Semantik

Diese Operation ist funktional identisch mit der Operation broadcast_in_dim, aber die Ergebnisform wird dynamisch über output_dimensions angegeben.

Der Vorgang akzeptiert auch die optionalen Attribute known_expanding_dimensions und known_nonexpanding_dimensions, um statisches Wissen über das Expansionsverhalten von Dimensionen auszudrücken. Wenn nichts angegeben ist, wird davon ausgegangen, dass alle Dimensionen möglicherweise erweitert werden.

Eingaben

Ausgaben

Einschränkungen

• (C1) element_type(result) wird so berechnet:
  • element_type(operand), wenn !is_per_axis_quantized(operand).
  • element_type(operand), mit der Ausnahme, dass quantization_dimension(operand), scales(operand) und zero_points(operand) von quantization_dimension(result), scales(result) und zero_points(result) abweichen können.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Für alle d in axes(operand):
  • dim(operand, d) = 1 oder
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Wenn is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Bei dim(operand, quantization_dimension(operand)) = 1, dann 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).

Beispiele

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

 Weitere Beispiele

dynamic_conv

Semantik

Diese Operation ist funktional identisch mit der Faltungsoperation, aber das Padding wird dynamisch über padding angegeben.

Eingaben