Indexierungsanalyse

Die HLO-Indexierungsanalyse ist eine Datenflussanalyse, die beschreibt, wie Elemente eines Tensors über „Indexierungskarten“ mit einem anderen Tensor in Beziehung stehen. Beispielsweise, wie die Indexe der Ausgabe einer HLO-Anweisung den Indexen der HLO-Anweisungsoberanden zugeordnet werden.

Beispiel

Für eine Broadcast-Nachricht von tensor<20xf32> an tensor<10x20x30xf32>

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

Die Indexierungszuordnung von der Ausgabe zur Eingabe ist (i, j, k) -> (j) für i in [0, 10], j in [0, 20] und k in [0, 30].

Motivation

XLA verwendet mehrere benutzerdefinierte Lösungen, um das Zusammenführen, die Operandennutzung und die Kachelungsschemas zu berücksichtigen (weitere Informationen unten). Ziel der Indexierungsanalyse ist es, eine wiederverwendbare Komponente für solche Anwendungsfälle bereitzustellen. Die Indexierungsanalyse basiert auf der Affine Map-Infrastruktur von MLIR und fügt HLO-Semantik hinzu.

Zusammenführen

Die Argumentation zur Speicherverschmelzung wird für nicht triviale Fälle möglich, wenn wir wissen, welche Elemente/Slices der Eingaben gelesen werden, um ein Element der Ausgabe zu berechnen.

Operandauslastung

Die Operandennutzung in XLA gibt an, wie stark jede Eingabe der Anweisung genutzt wird, wenn die Ausgabe vollständig genutzt wird. Derzeit wird die Auslastung auch nicht für einen generischen Fall berechnet. Mithilfe der Indexierungsanalyse können wir die Auslastung genau berechnen.

Kachelung:

Eine Kachel/ein Slice ist eine hyperrechteckige Teilmenge eines Tensors, die durch Offsets, Größen und Strides parametrisiert wird. Die Kachelweitergabe ist eine Möglichkeit, die Kachelparameter des Erstellers/Verbrauchers des Vorgangs anhand der Kachelparameter des Vorgangs selbst zu berechnen. Es gibt bereits eine Bibliothek, die das für Softmax und Dot erledigt. Die Kachelweitergabe kann allgemeiner und robuster gestaltet werden, wenn sie über Indexkarten ausgedrückt wird.

Indexierungskarte

Eine Indexierungsübersicht ist eine Kombination aus

eine symbolisch ausgedrückte Funktion, die jedes Element eines Tensors A auf Bereiche von Elementen im Tensor B abbildet;
Einschränkungen für gültige Funktionsargumente, einschließlich des Definitionsbereichs der Funktion.

Funktionsargumente werden in drei Kategorien unterteilt, um ihre Art besser zu kommunizieren:

Dimensionsvariablen des Tensors A oder eines GPU-Rasters, aus dem wir zuordnen; Werte sind statisch bekannt. Indexelemente werden auch als Dimensionsvariablen bezeichnet.
Bereichsvariablen. Sie definieren eine 1:n-Zuordnung und geben eine Reihe von Elementen in B an, die zum Berechnen eines einzelnen Werts von A verwendet werden. Die Werte sind statisch bekannt. Die kontrahierende Dimension einer Matrixmultiplikation ist ein Beispiel für eine Bereichsvariable.
Laufzeitvariablen, die nur während der Ausführung bekannt sind. Beispiel: Das Argument „indices“ des Vorgangs gather.

Das Ergebnis der Funktion ist ein Index des Ziel-B-Tensors.

Kurz gesagt: Eine Indexierungsfunktion von Tensor A zu Tensor B für den Vorgang x ist

map_ab(index in A, range variables, runtime variables) -> index in B.

Um die verschiedenen Arten von Mapping-Argumenten besser zu unterscheiden, schreiben wir sie so:

map_ab(index in A)[range variables]{runtime variables} -> (index in B)

Sehen wir uns beispielsweise die Indexierungszuordnungen für den Reduce-Vorgang f32[4, 8] out = reduce(f32[2, 4, 8, 16] in, 0), dimensions={0,3} an:

Um Elemente von in auf out abzubilden, kann unsere Funktion als (d0, d1, d2, d3) -> (d1, d2) ausgedrückt werden. Die Einschränkungen der Variablen d0 in [0, 1], d1 in [0, 3], d2 in [0, 7], d3 in [0, 15] werden durch die Form von in definiert.
Elemente von out auf in abbilden: out hat nur zwei Dimensionen und durch die Reduzierung werden zwei Bereichsvariablen eingeführt, die die reduzierenden Dimensionen abdecken. Die Zuordnungsfunktion ist also (d0, d1)[s0, s1] -> (s0, d0, d1, s1), wobei (d0, d1) der Index von out ist. s0 und s1 sind Bereiche, die durch die Semantik des Vorgangs definiert werden und die Dimensionen 0 und 3 des in-Tensors umfassen. Die Einschränkungen sind d0 in [0, 3], d1 in [0, 7], s0 in [0,1], s1 in [0, 15].

In den meisten Fällen geht es darum, die Elemente der Ausgabe zuzuordnen. Für Berechnungen

C = op1(A, B)
E = op2(C, D)

Wir können von „Indexierung von B“ sprechen, was „Zuordnung von Elementen von E zu den Elementen von B“ bedeutet. Das kann im Vergleich zu anderen Arten der Datenflussanalyse, die von der Eingabe zur Ausgabe verlaufen, kontraintuitiv sein.

Einschränkungen für Variablen ermöglichen Optimierungschancen und helfen bei der Codegenerierung. In der Dokumentation und bei den Implementierungseinschränkungen wird auch auf domain verwiesen, da sie alle gültigen Kombinationen oder Argumentwerte der Zuordnungsfunktion definieren. Bei vielen Operationen beschreiben Einschränkungen einfach die Dimensionen von Tensoren. Bei einigen Operationen können sie jedoch komplizierter sein. Siehe Beispiele unten.

Da Funktionen und Argumentbeschränkungen symbolisch ausgedrückt werden und Funktionen und Beschränkungen kombiniert werden können, lässt sich eine kompakte Indexierungszuordnung für eine beliebig große Berechnung (Fusion) berechnen.

Die Ausdrucksstärke der symbolischen Funktion und der Einschränkungen ist ein Kompromiss zwischen der Komplexität der Implementierung und den Optimierungsvorteilen, die sich aus einer präziseren Darstellung ergeben. Bei einigen HLO-Vorgängen werden Zugriffsmuster nur ungefähr erfasst.

Implementierung

Da wir die Anzahl der Neuberechnungen minimieren möchten, benötigen wir eine Bibliothek für symbolische Berechnungen. XLA basiert bereits auf MLIR. Daher verwenden wir mlir::AffineMap, anstatt eine weitere symbolische Arithmetikbibliothek zu schreiben.

Ein typisches AffineMap sieht so aus:

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

AffineMap hat zwei Arten von Parametern: Dimensionen und Symbole. Dimensionen entsprechen den Dimensionsvariablen d, Symbole den Bereichsvariablen r und den Laufzeitvariablen rt. AffineMap enthält keine Metadaten zu Einschränkungen der Parameter. Diese müssen also separat angegeben werden.

struct Interval {
 int64_t lower;
 int64_t upper;
};

class IndexingMap {
   // Variable represents dimension, range or runtime variable.
  struct Variable {
    Interval bounds;
    // Name of the variable is used for nicer printing.
    std::string name = "";
  };

  mlir::AffineMap affine_map_;

  // DimVars represent dimensions of a tensor or of a GPU grid.
  std::vector<Variable> dim_vars_;

  // RangeVars represent ranges of values, e.g. to compute a single element of
  // the reduction's result we need a range of values from the input tensor.
  std::vector<Variable> range_vars_;

  // RTVars represent runtime values, e.g. a dynamic offset in
  // HLO dynamic-update-slice op.
  std::vector<Variable> rt_vars_;
  llvm::DenseMap<mlir::AffineExpr, Interval> constraints_;
};

dim_vars_ codiert die inklusiven Box-Einschränkungen für die Dimensionsvariablen d der Indexierungszuordnung, die in der Regel mit der Form des Ausgabetensors für Vorgänge wie „transpose“, „reduce“, „elementwise“ und „dot“ übereinstimmen. Es gibt jedoch einige Ausnahmen wie HloConcatenateInstruction.

range_vars_ – alle Werte, die Bereichsvariablen s annehmen. Die Bereichsvariablen sind erforderlich, wenn mehrere Werte zum Berechnen eines einzelnen Elements des Tensors benötigt werden, von dem wir die Zuordnung vornehmen, z.B. für die Indexzuordnung von Ausgaben zu Eingaben bei Reduzierungen oder die Zuordnung von Eingaben zu Ausgaben bei Broadcasts.

rt_vars_ codiert die zulässigen Werte zur Laufzeit. Der Offset ist beispielsweise für eine eindimensionale HloDynamicSliceInstruction dynamisch. Der entsprechende RTVar-Wert liegt zwischen 0 und tensor_size - slice_size - 1.

constraints_ erfasst Beziehungen zwischen Werten im Format <expression> in <range>, z.B. d0 + s0 in [0, 20]. Zusammen mit Variable.bounds definieren sie die „Domain“ der Indexierungsfunktion.

Sehen wir uns anhand von Beispielen an, was das alles bedeutet.

Karten für nicht zusammengeführte Vorgänge indexieren

Elementweise

Bei elementweisen Operationen ist die Indexierungszuordnung eine Identität.

  p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  output = f32[10, 20] add(p0, p1)

Die Zuordnung von Ausgabe zu Eingabe output -> p0:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

Die Zuordnung von Eingabe zu Ausgabe p0 -> output:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

Nachricht an alle

Beim Broadcasting werden einige Dimensionen entfernt, wenn wir die Ausgabe der Eingabe zuordnen, und hinzugefügt, wenn wir die Eingabe der Ausgabe zuordnen.

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

Die Zuordnung von Ausgabe zu Eingabe:

(d0, d1, d2) -> (d1),
domain:
d0 in [0, 9],
d1 in [0, 19],
d2 in [0, 29]

Die Zuordnung von Eingabe zu Ausgabe:

(d0)[s0, s1] -> (s0, d0, s1),
domain:
d0 in [0, 19],
s0 in [0, 9],
s1 in [0, 29]

Beachten Sie, dass wir jetzt Bereichsvariablen s auf der rechten Seite für die Zuordnung von Eingabe zu Ausgabe haben. Das sind die Symbole, die Wertebereiche darstellen. In diesem speziellen Fall wird beispielsweise jedes Element der Eingabe mit dem Index d0 einem 10 × 1 × 30-Slice der Ausgabe zugeordnet.

Iota

„Iota“ hat keinen Eingabetensor-Operanden, daher gibt es keine Eingabeindex-Argumente.

iota = f32[2,4] iota(), dimensions={1}

Ausgabe-zu-Eingabe-Zuordnung:

(d0, d1) -> ()
domain:
d0 in [0, 1]
d1 in [0, 3]

Eingabe-Ausgabe-Zuordnung:

()[s0, s1] -> (s0, s1)
domain:
s0 in [0, 1]
s1 in [0, 3]

DynamicSlice

DynamicSlice hat Offsets, die nur zur Laufzeit bekannt sind.

src = s32[2, 2, 258] parameter(0)
of1 = s32[] parameter(1)
of2 = s32[] parameter(2)
of3 = s32[] parameter(3)
ds = s32[1, 2, 32] dynamic-slice(src, of1, of2, of3), dynamic_slice_sizes={1, 2, 32}

Die Zuordnung von Ausgabe zu Eingabe von ds zu src:

(d0, d1, d2){rt0, rt1, rt2} -> (d0 + rt0, d1 + rt1, d2 + rt2),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31],
rt0 in [0, 1],
rt1 in [0, 0],
rt2 in [0, 226]

Beachten Sie, dass wir jetzt rt auf der rechten Seite für die Zuordnung von Ein- zu Ausgaben haben. Das sind die Symbole, die Laufzeitwerte darstellen. In diesem speziellen Fall greifen wir beispielsweise für jedes Element der Ausgabe mit den Indexwerten d0, d1, d2 auf die Slice-Offsets of1, of2 und of3 zu, um den Index der Eingabe zu berechnen. Die Intervalle für die Laufzeitvariablen werden abgeleitet, indem davon ausgegangen wird, dass der gesamte Slice innerhalb der Grenzen bleibt.

Die Ausgabe-zu-Eingabe-Zuordnung für of1, of2 und of3:

(d0, d1, d2) -> (),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31]

DynamicUpdateSlice

src = s32[20,30] parameter(0)
upd = s32[5,10] parameter(1)
of1 = s32[] parameter(2)
of2 = s32[] parameter(3)
dus = s32[20,30] dynamic-update-slice(
    s32[20,30] src, s32[5,10] upd, s32[] of1, s32[] of2)

Die Zuordnung von Ausgabe zu Eingabe für src ist trivial. Sie kann präziser gemacht werden, indem die Domain auf die nicht aktualisierten Indexe beschränkt wird. Ungleichheitsbedingungen werden in Indexierungszuordnungen jedoch derzeit nicht unterstützt.

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 19],
d1 in [0, 29]

Die Zuordnung von Ausgabe zu Eingabe für upd:

(d0, d1){rt0, rt1} -> (d0 - rt0, d1 - rt1),
domain:
d0 in [0, 19],
d1 in [0, 29],
rt0 in [0, 15],
rt1 in [0, 20]

Beachten Sie, dass wir jetzt rt0 und rt1 haben, die Laufzeitwerte darstellen. In diesem speziellen Fall greifen wir für jedes Element der Ausgabe mit den Indexen d0, d1 auf die Slice-Offsets of1 und of2 zu, um den Index der Eingabe zu berechnen. Die Intervalle für die Laufzeitvariablen werden abgeleitet, indem davon ausgegangen wird, dass der gesamte Slice innerhalb der Grenzen bleibt.

Die Zuordnung von Ausgabe zu Eingabe für of1 und of2:

(d0, d1) -> (),
domain:
d0 in [0, 19],
d1 in [0, 29]

Gather

Es wird nur die vereinfachte Erfassung unterstützt. Weitere Informationen finden Sie unter gather_simplifier.h.

operand = f32[33,76,70] parameter(0)
indices = s32[1806,2] parameter(1)
gather = f32[1806,7,8,4] gather(operand, indices),
  offset_dims={1,2,3},
  collapsed_slice_dims={},
  start_index_map={0,1},
  index_vector_dim=1,
  slice_sizes={7,8,4}

Die Zuordnung von Ausgabe zu Eingabe für operand:

(d0, d1, d2, d3){rt0, rt1} -> (d1 + rt0, d2 + rt1, d3),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
rt0 in [0, 26],
rt1 in [0, 68]

Beachten Sie, dass wir jetzt rt-Symbole haben, die Laufzeitwerte darstellen.

Die Zuordnung von Ausgabe zu Eingabe für indices:

(d0, d1, d2, d3)[s0] -> (d0, s0),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
s0 in [0, 1]

Die Bereichsvariable s0 zeigt, dass wir die gesamte Zeile (d0, *) des indices-Tensors benötigen, um ein Element der Ausgabe zu berechnen.

Achsen tauschen

Die Indexierungszuordnung für die Transponierung ist eine Permutation der Ein-/Ausgabedimensionen.

p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}

Die Zuordnung von Ausgabe zu Eingabe:

(d0, d1, d2, d3) -> (d0, d3, d1, d2),
domain:
d0 in [0, 2],
d1 in [0, 5],
d2 in [0, 127],
d3 in [0, 12287],

Die Zuordnung von Eingabe zu Ausgabe:

(d0, d1, d2, d3) -> (d0, d2, d3, d1),
domain:
d0 in [0, 2],
d1 in [0, 12287],
d2 in [0, 5],
d3 in [0, 127]

Umkehren

Die Indexierungszuordnung für die Rückgängigmachung ändert die zurückgesetzten Dimensionen in upper_bound(d_i) - d_i:

p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}

Die Zuordnung von Ausgabe zu Eingabe:

(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3),
domain:
d0 in [0, 0],
d1 in [0, 16],
d2 in [0, 8],
d3 in [0, 8]

Die Zuordnung von Eingabe zu Ausgabe:

(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3),
domain:
d0 in [0, 0],
d1 in [0, 16],
d2 in [0, 8],
d3 in [0, 8]

(Variadic)Reduce

Variadische Reduzierung hat mehrere Eingaben und mehrere Anfangswerte. Die Zuordnung von Ausgabe zu Eingabe fügt die reduzierten Dimensionen hinzu.

p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
out = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
  dimensions={0}, to_apply=max

Die Ausgabe-zu-Eingabe-Zuordnungen:

out[0] -> p0:

(d0)[s0] -> (s0, d0),
domain:
d0 in [0, 9],
s0 in [0, 255]

out[0] -> p0_init:

(d0) -> (),
domain:
d0 in [0, 9]

Die Zuordnungen von Eingaben zu Ausgaben:

p0 -> out[0]:

(d0, d1) -> (d1),
domain:
d0 in [0, 255],
d1 in [0, 9]

p0_init -> out[0]:

()[s0] -> (s0),
domain:
s0 in [0, 9]

Slice

Die Indexierung von der Ausgabe zur Eingabe für den Slice führt zu einer Indexierungszuordnung mit Schrittweite, die für jedes Element der Ausgabe gültig ist. Die Zuordnung von der Eingabe zur Ausgabe ist auf einen gestuften Bereich der Elemente in der Eingabe beschränkt.

p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
  slice={[5:10:1], [3:20:7], [0:50:2]}

Die Zuordnung von Ausgabe zu Eingabe:

(d0, d1, d2) -> (d0 + 5, d1 * 7 + 3, d2 * 2),
domain:
d0 in [0, 4],
d1 in [0, 2],
d2 in [0, 24]

Die Zuordnung von Eingabe zu Ausgabe:

(d0, d1, d2) -> (d0 - 5, (d1 - 3) floordiv 7, d2 floordiv 2),
domain:
d0 in [5, 9],
d1 in [3, 17],
d2 in [0, 48],
(d1 - 3) mod 7 in [0, 0],
d2 mod 2 in [0, 0]

Umformen

Es gibt verschiedene Arten von Reshapes.

Form minimieren

Dies ist eine „linearisierende“ Umformung von N-D zu 1D.

p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)

Die Zuordnung von Ausgabe zu Eingabe:

(d0) -> (d0 floordiv 8, d0 mod 8),
domain:
d0 in [0, 31]

Die Zuordnung von Eingabe zu Ausgabe:

(d0, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

Form maximieren

Dies ist eine inverse „collapse shape“-Operation, bei der eine 1D-Eingabe in eine N-D-Ausgabe umgeformt wird.

p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)

Die Zuordnung von Ausgabe zu Eingabe:

(d0, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

Die Zuordnung von Eingabe zu Ausgabe:

(d0) -> (d0 floordiv 8, d0 mod 8),
domain:
d0 in [0, 31]

Generisches Umformen

Dies sind die Reshape-Vorgänge, die nicht als einzelne Expand- oder Collapse-Form dargestellt werden können. Sie können nur als Kombination aus zwei oder mehr Formen zum Ein- oder Ausblenden dargestellt werden.

Beispiel 1: Linearisierung und Delinearisierung

p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)

Diese Umformung kann als Zusammensetzung der Reduzierung der Form von tensor<4x8xf32> auf tensor<32xf32> und dann einer Formexpansion auf tensor<2x4x4xf32> dargestellt werden.

Die Zuordnung von Ausgabe zu Eingabe:

(d0, d1, d2) -> (d0 * 2 + d1 floordiv 2, d2 + (d1 mod 2) * 4),
domain:
d0 in [0, 1],
d1 in [0, 3],
d2 in [0, 3]

Die Zuordnung von Eingabe zu Ausgabe:

(d0, d1) -> (d0 floordiv 2, d1 floordiv 4 + (d0 mod 2) * 2, d1 mod 4),
domain:
d0 in [0, 3],
d1 in [0, 7]

Beispiel 2: Erweiterte und minimierte untergeordnete Formen

p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)

Diese Umformung kann als Zusammensetzung von zwei Umformungen dargestellt werden. Bei der ersten werden die äußersten Dimensionen tensor<4x8x12xf32> auf tensor<32x12xf32> reduziert und bei der zweiten wird die innerste Dimension tensor<32x12xf32> auf tensor<32x3x4xf32> erweitert.

Die Zuordnung von Ausgabe zu Eingabe:

(d0, d1, d2) -> (d0 floordiv 8, d0 mod 8, d1 * 4 + d2),
domain:
d0 in [0, 31],
d1 in [0, 2]
d2 in [0, 3]

Die Zuordnung von Eingabe zu Ausgabe:

(d0, d1, d2) -> (d0 * 8 + d1, d2 floordiv 4, d2 mod 4),
domain:
d0 in [0, 3],
d1 in [0, 7],
d2 in [0, 11]

Bitcast

Ein Bitcast-Vorgang kann als Folge von Transponieren, Umformen und Transponieren dargestellt werden. Daher sind die Indexierungskarten nur eine Zusammensetzung der Indexierungskarten für diese Sequenz.

Concatenate

Die Ausgabe-zu-Eingabe-Zuordnung für „concat“ ist für alle Eingaben definiert, aber mit nicht überlappenden Bereichen. Das bedeutet, dass jeweils nur eine der Eingaben verwendet wird.

p0 = f32[2, 5, 7] parameter(0)
p1 = f32[2, 11, 7] parameter(1)
p2 = f32[2, 17, 7] parameter(2)
ROOT output = f32[2, 33, 7] concatenate(f32[2, 5, 7] p0, f32[2, 11, 7] p1, f32[2, 17, 7] p2), dimensions={1}

Die Ausgabe wird den Eingaben zugeordnet:

output -> p0:

(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 4],
d2 in [0, 6]

output -> p1:

(d0, d1, d2) -> (d0, d1 - 5, d2),
domain:
d0 in [0, 1],
d1 in [5, 15],
d2 in [0, 6]

output -> p2:

(d0, d1, d2) -> (d0, d1 - 16, d2),
domain:
d0 in [0, 1],
d1 in [16, 32],
d2 in [0, 6]

Die Eingaben für Ausgabekarten:

p0 -> output:

(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 4],
d2 in [0, 6]

p1 -> output:

(d0, d1, d2) -> (d0, d1 + 5, d2),
domain:
d0 in [0, 1],
d1 in [0, 10],
d2 in [0, 6]

p2 -> output:

(d0, d1, d2) -> (d0, d1 + 16, d2),
domain:
d0 in [0, 1],
d1 in [0, 16],
d2 in [0, 6]

Punkt

Die Indexierung von Maps für „dot“ ähnelt sehr der von „reduce“.

p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
output = f32[4, 128, 64] dot(p0, p1),
  lhs_batch_dims={0}, rhs_batch_dims={0},
  lhs_contracting_dims={2}, rhs_contracting_dims={1}

Die Ausgabe wird den Eingaben zugeordnet:

Ausgabe -> p0:

(d0, d1, d2)[s0] -> (d0, d1, s0),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 63],
s0 in [0, 255]

Ausgabe -> p1:

(d0, d1, d2)[s0] -> (d0, s0, d2),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 63],
s0 in [0, 255]

Die Eingaben für Ausgabekarten:

p0 -> output:

(d0, d1, d2)[s0] -> (d0, d1, s0),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 255],
s0 in [0, 63]

p1 -> Ausgabe:

(d0, d1, d2)[s0] -> (d0, s0, d1),
domain:
d0 in [0, 3],
d1 in [0, 255],
d2 in [0, 63],
s0 in [0, 127]

Pad

Die Indexierung von PadOp ist das Gegenteil der Indexierung von SliceOp.

p0 = f32[4, 4] parameter(0)
p1 = f32[] parameter(1)
pad = f32[12, 16] pad(p0, p1), padding=1_4_1x4_8_0

Die Padding-Konfiguration 1_4_1x4_8_0 steht für lowPad_highPad_interiorPad_dim_0 x lowPad_highPad_interiorPad_dim_1.

Die Ausgabe-zu-Eingabe-Zuordnungen:

Ausgabe -> p0:

(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4),
domain:
d0 in [1, 7],
d1 in [4, 7],
(d0 - 1) mod 2 in [0, 0]

Ausgabe -> p1:

(d0, d1) -> (),
domain:
d0 in [0, 11],
d1 in [0, 15]

ReduceWindow

ReduceWindow in XLA führt auch Padding durch. Daher können die Indexierungszuordnungen als Zusammensetzung der ReduceWindow-Indexierung ohne Padding und der PadOp-Indexierung berechnet werden.

c_inf = f32[] constant(-inf)
p0 = f32[1024, 514] parameter(0)
outpu = f32[1024, 3] reduce-window(p0, c_inf),
  window={size=1x512 pad=0_0x0_0}, to_apply=max

Die Ausgabe-zu-Eingabe-Zuordnungen:

output -> p0:

(d0, d1)[s0] -> (d0, d1 + s0),
domain:
d0 in [0, 1023],
d1 in [0, 2],
s0 in [0, 511]

output -> c_inf:

(d0, d1) -> (),
domain:
d0 in [0, 1023],
d1 in [0, 2]

Karten für Fusion indexieren

Die Indexierungszuordnung für den Fusionsvorgang ist eine Zusammensetzung der Indexierungszuordnungen für jeden Vorgang im Cluster. Es kann vorkommen, dass einige Eingaben mehrmals mit unterschiedlichen Zugriffsmustern gelesen werden.

Eine Eingabe, mehrere Indexierungskarten

Hier ist ein Beispiel für p0 + transpose(p0).

f {
  p0 = f32[1000, 1000] parameter(0)
  transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
  ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}

Die Zuordnungen von Ausgabe zu Eingabe für p0 sind (d0, d1) -> (d0, d1) und (d0, d1) -> (d1, d0). Das bedeutet, dass wir den Eingabeparameter möglicherweise zweimal lesen müssen, um ein Element der Ausgabe zu berechnen.

Eine Eingabe, deduplizierte Indexierungszuordnung

Es gibt Fälle, in denen die Indexierungskarten tatsächlich identisch sind, auch wenn das nicht sofort ersichtlich ist.

f {
  p0 = f32[20, 10, 50] parameter(0)
  lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
  lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
  lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
  rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
  rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
  rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
  ROOT output = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}

Die Indexzuordnung von Ausgabe zu Eingabe für p0 ist in diesem Fall einfach (d0, d1, d2) -> (d2, d0, d1).

Softmax

Die Zuordnungen von Ausgabe- zu Eingabeindex für parameter 0 für Softmax:

(d0, d1, d2)[s0] -> (d0, d1, s0),
domain:
d0 in [0, 1],
d1 in [0, 64],
d2 in [0, 124],
s0 in [0, 124]

und

(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 64],
d2 in [0, 124]

Dabei bezieht sich s0 auf die innerste Dimension der Eingabe.

Weitere Beispiele finden Sie unter indexing_analysis_test.cc.

Tool zur Indexierung von Karten

Der Standard-Simplifier für mlir::AffineMap Upstream kann keine Annahmen über die Bereiche von Dimensionen/Symbolen treffen. Daher können Ausdrücke mit mod und div nicht effizient vereinfacht werden.

Wir können das Wissen über die Unter- und Obergrenzen der Unterausdrücke in den affinen Abbildungen nutzen, um sie noch weiter zu vereinfachen.

Der Simplifier kann die folgenden Ausdrücke neu schreiben.

(d0, d1) -> (d0 + d1 floordiv 16, d1 mod 16) für d in [0, 6] x [0, 14] wird zu (d0, d1) -> (d0, d1)
(d0, d1, d2) -> ((100d0 + 10d1 + d2) floorDiv 100, ((100d0 + 10d1 + d2) mod 100) floordiv 10, d2 mod 10) für di in [0, 9] wird zu (d0, d1, d2) -> (d0, d1, d2).
(d0, d1, d2) -> ((16d0 + 4d1 + d2) floordiv 8, (16d0 + 4d1 + d2) mod 8) für d_i in [0, 9] wird zu (d0, d1, d2) -> (2d0 + (4d1 + d2) floordiv 8,(4d1 + d2) mod 8).
(d0, d1) -> (-(-11d0 - d1 + 109) floordiv 11 + 9) für d in [0, 9] x [0, 10] wird zu (d0, d1) -> (d0).

Mit dem Indexierungs-Map-Vereinfacher können wir erkennen, dass sich einige der verketteten Reshape-Vorgänge in HLO gegenseitig aufheben.

p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)

Nach der Erstellung von Indexkarten und ihrer Vereinfachung erhalten wir

(d0, d1, d2) -> (d0, d1, d2).

Durch die Vereinfachung der Indexierungskarte werden auch die Einschränkungen vereinfacht.

Einschränkungen vom Typ lower_bound <= affine_expr (floordiv, +, -, *) constant <= upper_bound werden als updated_lower_bound <= affine_expr <= updated_upped_bound neu geschrieben.
Bedingungen, die immer erfüllt sind, z.B. d0 + s0 in [0, 20] für d0 in [0, 5] und s0 in [1, 3], werden entfernt.
Affine Ausdrücke in den Einschränkungen werden wie die affine Indexierungsabbildung oben optimiert.

Weitere Beispiele finden Sie unter indexing_map_test.cc.

Indexierungsanalyse Mit Sammlungen den Überblick behalten Sie können Inhalte basierend auf Ihren Einstellungen speichern und kategorisieren.

Beispiel

Motivation

Zusammenführen

Operandauslastung

Kachelung:

Indexierungskarte

Implementierung

Karten für nicht zusammengeführte Vorgänge indexieren

Elementweise

Form minimieren

Form maximieren

Generisches Umformen

Beispiel 1: Linearisierung und Delinearisierung

Beispiel 2: Erweiterte und minimierte untergeordnete Formen

Bitcast

Karten für Fusion indexieren

Eine Eingabe, mehrere Indexierungskarten

Eine Eingabe, deduplizierte Indexierungszuordnung

Softmax

Tool zur Indexierung von Karten

Indexierungsanalyse