In diesem Dokument wird die HLO-Indexierungsanalyse beschrieben, mit der Sie Indexierungskarten für HLO-Vorgänge symbolisch berechnen können. Die Indexierungszuordnung ist eine Funktion, die Indizes eines Tensors den Indizes eines anderen zuordnet, z.B. die Indizes einer HLO-Anweisungsausgabe den Indizes der HLO-Anweisungseingaben oder umgekehrt.
Beispiel
Für eine Übertragung von tensor<20xf32> bis tensor<10x20x30xf32>
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
Die Indexierungszuordnung von der Ausgabe in die Eingabe ist \((i, j, k) \mapsto (j)\) for $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$.
Ziel
XLA GPU verwendet mehrere maßgeschneiderte Lösungen zur Logik von Zusammenführung, Operandenauslastung und Tiling-Schemas (weitere Details unten). Das Ziel der Indexierungsanalyse besteht darin, eine wiederverwendbare Komponente für solche Anwendungsfälle bereitzustellen. Die Indexierungsanalyse basiert auf der Affine Map-Infrastruktur von MLIR und verfügt über eine HLO-Semantik.
Koaleszen
Die Logik der Speicherzusammenführung wird auch bei nicht trivialen Fällen machbar, wenn wir wissen, welche Elemente/Slices der Eingaben gelesen werden, um ein Element der Ausgabe zu berechnen.
Operand-Auslastung
Die Operandenauslastung in XLA gibt an, wie viel jede Eingabe der Anweisung verwendet wird, vorausgesetzt, ihre Ausgabe ist vollständig ausgelastet. Derzeit wird die Auslastung auch nicht für einen generischen Fall berechnet. Die Indexanalyse ermöglicht eine präzise Berechnung der Auslastung.
Kachelung:
Eine Kachel/ein Segment ist eine hyperrechteckige Teilmenge eines Tensors, die durch Versätze, Größen und Schritte parametriert ist. Die Kachelweitergabe ist eine Möglichkeit, die Kachelparameter des Erstellers/Nutzers des Vorgangs mithilfe der Tiling-Parameter des Vorgangs selbst zu berechnen. Es gibt bereits eine Bibliothek, die dies für Softmax und Punkt erledigt. Die Verbreitung von Kacheln kann allgemeiner und robuster gestaltet werden, wenn dies über Indexierungskarten erfolgt.
Funktion und Domain
Die Indexierungszuordnung ist eine Funktion, \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\)die einen Multiindex \(\boldsymbol{d}\) eines Tensors \(A\) den Elementen/Bereichen des Tensors \(B\)zuordnet. Der Parameter \(\boldsymbol{s}\) bezieht sich auf die Indexbereiche der Dimensionen, die in Tensor \(B\), aber nicht in Tensor \(A\)vorhanden sind.
Wenn wir beispielsweise eine Reduktion von tensor<2x4x8x16xf32> auf tensor<4x8xf32> haben, lautet die Indexierungszuordnung von der 2D-Ausgabe zur 4D-Eingabe\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\), wobei \(d_i\) die Dimensionsparameter sind, die den Indizes des Ausgabetensors entsprechen. Parameter \(s_j\)codieren mehrere Werte. Um beispielsweise ein \((d_0, d_1)\) Element der Ausgabe zu berechnen, benötigen wir \((s_0, d_0, d_1, s_1)\) Elemente der Eingabe, wobei \(s_0 \in [0, 2)\) und\(s_1 \in [0, 16)\).
Diese Zuordnung kann aus den Attributen von HLO-Anweisungen erstellt werden oder die Zuordnungen nicht zusammengeführter Anweisungen können für die Indexierung einer Fusion erstellt werden. Die Zuordnung hat auch eine Domain, die angibt, für welche Elemente des Tensors die Zuordnung existiert.
\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]
Da wir die Neuberechnung minimieren möchten, benötigen wir eine Bibliothek für symbolische Berechnungen. Da XLA bereits von MLIR abhängig ist, verwenden wir mlir::AffineMap, anstatt eine symbolische arithmetische Bibliothek zu schreiben.
Ein typisches AffineMap-Objekt sieht so aus
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
Bei AffineMap gibt es zwei praktische Parametertypen: Dimensionen und Symbole, die wir für \(\boldsymbol d\) und \(\boldsymbol s\) verwenden können.
AffineMap enthält keine Metadaten zu Dimensionsbereichen. Wir müssen diese Daten also selbst bereitstellen.
struct Range {
int64_t lower_bound;
int64_t upper_bound;
};
struct IndexingMap {
mlir::AffineMap affine_map;
std::vector<Range> dimension_ranges;
std::vector<Range> symbol_ranges;
llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};
dim_ranges codiert die inklusiven Boxeneinschränkungen für die Dimensionsparameter \(\boldsymbol{d}\) der Indexierungskarte, die normalerweise mit der Form des Ausgabetensors für Operationen wie Transponieren, Reduzieren, elementweise oder Punkt übereinstimmen. Es gibt jedoch einige Ausnahmen wie HloConcatenateInstruction.
symbol_ranges-codieren mögliche Werte, die \(\boldsymbol {s}\) Parameter annehmen können.
Sehen wir uns anhand eines Beispiels an, was das alles genau bedeutet.
Indexierung von Karten für nicht zusammengeführte Vorgänge
Elementweise
Bei elementweisen Vorgängen ist die Indexierungskarte eine Identität.
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
Die Ausgabe für die Eingabe entspricht:
- Ausgabe -> Eingabe_0: \((d_0, d_1) \mapsto (d_0, d_1)\) für $\boldsymbol{d} \in [0,9] \times [0, 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}(output)$
- Ausgabe -> Eingabe_1: \((d_0, d_1) \mapsto (d_0, d_1)\) für $\boldsymbol{d} \in {\rm Dom} (Ausgabe)$
Die Eingabe für Ausgabezuordnungen
- Eingabe_i -> Ausgabe: \((d_0, d_1) \mapsto (d_0, d_1)\) für $\boldsymbol{d} \in {\rm Dom}(output)$
Nachricht an alle
Beim Broadcasting werden einige Dimensionen entfernt, wenn die Ausgabe der Eingabe zugeordnet und bei der Zuordnung der Eingabe zur Ausgabe hinzugefügt wird.
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
Die Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe: \((d_0, d_1, d_2) \mapsto (d_1)\) for $\boldsymbol{d} \in {\rm Dom}(output)$
Die Eingabe für die Ausgabezuordnung
- Eingabe -> Ausgabe: \((d_0) \mapsto (s_0, d_1, s_1)\) für $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9] \times [0, 29]$.
Beachten Sie, dass sich jetzt \(\boldsymbol s\) auf der rechten Seite für die Eingabe/Ausgabe-Zuordnung befindet. Dies sind die Symbole, die Wertebereiche darstellen. In diesem speziellen Fall wird beispielsweise jedes Eingabeelement mit dem Index \(d_0\) einem Slice der Ausgabe von 10 × 1 × 30 zugeordnet.
Konstante und Iota
Praktischerweise haben sie keine Eingabeparameter, sodass für die Indexierung nichts berechnet werden muss.
Transponieren
Die Indexierungskarte für das Transponieren ist eine Permutation von Eingabe-/Ausgabedimensionen.
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
Die Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\) für\(\boldsymbol{d} \in {\rm Dom}(output)\)
Die Eingabe für die Ausgabezuordnung:
- Eingabe -> Ausgabe: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\) für\(\boldsymbol{d} \in {\rm Dom}(input)\)
Reverse
Die Indexierungskarte für umgekehrt ä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 Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(output)$
Die Eingabe für die Ausgabezuordnung:
- Eingabe -> Ausgabe: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(input)$
(Variadic)Einschränken
Die variative Reduktion hat mehrere Eingaben und mehrere Initialisierungen. Bei der Zuordnung von der Ausgabe zur Eingabe werden die reduzierten Dimensionen hinzugefügt. Sie verhält sich also in gewisser Weise wie eine Umkehrung zu einem Broadcast.
p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
dimensions={0}, to_apply=min
Die Ausgabe für die Eingabe entspricht:
- Ausgabe -> input_j: \((d_0) \mapsto (s_0, d_0)\) for $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9]$
- Ausgabe -> init_j: \((d_0) \mapsto ()\) für $\boldsymbol{d} \in {\rm Dom}(output)$
Die Eingabe für die Ausgabe ordnet Folgendes zu:
- Eingabe_i -> Ausgabe_j: \((d_0, d_1) \mapsto (d_1)\) für $\boldsymbol{d} \in {\rm Dom}(input)$
- init_i ->output_j: \(() \mapsto (s_0)\) für \(\boldsymbol{s} \in [0, 9]\)
für \(i, j = 0, \ldots, INPUT\\_COUNT\).
Slice
Die Indexierung von Ausgabe zu Eingabe für Slice führt zu einer gestaffelten Indexierungskarte, die für jedes Element der Ausgabe gültig ist. Die Zuordnung von Eingabe zu Ausgabe ist auf einen abgegrenzten 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 Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe: \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\) für\(\boldsymbol{d} \in {\rm Dom}(output)\)
Die Eingabe für die Ausgabezuordnung:
- Eingabe -> Ausgabe: \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) für\(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) mit den Schritten $[1, 7, 2]$.
Noch offen: Indexierung von Eingabe/Ausgabe
Form ändern
Reshapes gibt es in verschiedenen Geschmacksrichtungen.
Form minimieren
Dies ist eine „linearisierte“ Umformung von N–D nach 1D.
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
Die Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) for $\boldsymbol{d} \in {\rm Dom}(output)$
Die Eingabe für die Ausgabezuordnung:
- Eingabe -> Ausgabe: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) für $\boldsymbol{d} \in {\rm Dom}(input)$.
Form maximieren
Dies ist eine inverse „Minimierungsform“-Operation, die eine 1D-Eingabe in eine N-D-Ausgabe umformt.
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
Die Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) for $\boldsymbol{d} \in {\rm Dom}(output)$
Die Eingabe für die Ausgabezuordnung:
- Eingabe -> Ausgabe: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) für $\boldsymbol{d} \in {\rm Dom}(input)$.
Allgemeine Formulierung
Dies sind die Vorgänge zur Umformung, die nicht als einzelne Maximierungs- oder Minimierungsform dargestellt werden können. Sie können nur als Zusammensetzung aus zwei oder mehr Erweiterungs- oder Minimierungsformen dargestellt werden.
Beispiel 1: Linearisierung/Delinearisierung.
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
Diese Umformung kann als Zusammensetzung der Minimierungsform von tensor<4x8xf32> zu tensor<32xf32> und dann als Erweiterung der Form auf tensor<2x4x4xf32> dargestellt werden.
Die Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$
für \(\boldsymbol{d} \in {\rm Dom}(output)\)
Die Eingabe für die Ausgabezuordnung:
- Eingabe -> Ausgabe: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$
für \(\boldsymbol{d} \in {\rm Dom}(input)\).
Beispiel 2: Maximierte und minimierte Unterformen
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
Diese Umformung kann als Kombination aus zwei Umformungen dargestellt werden. Die erste minimiert die äußersten Dimensionen tensor<4x8x12xf32> auf tensor<32x12xf32> und die zweite erweitert die innerste Dimension tensor<32x12xf32> auf tensor<32x3x4xf32>.
Die Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe: \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\) für \(\boldsymbol{d} \in {\rm Dom}(output)\)
Die Eingabe für die Ausgabezuordnung:
- Eingabe -> Ausgabe: \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\)für \(\boldsymbol{d} \in {\rm Dom}(input)\).
Bitcast
Ein Bitcast-Vorgang kann als Sequenz von Transpose-Reshape-Transpose dargestellt werden. Daher sind die Indexierungskarten nur eine Zusammensetzung von Indexierungskarten für diese Sequenz.
Verketten
Die Ausgabe-zu-Eingabe-Zuordnung für Concat ist für alle Eingaben definiert, aber mit nicht überlappenden Domains, d.h., es wird immer nur eine der Eingaben verwendet.
p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
dimensions={1}
Die Ausgabe für die Eingabezuordnung:
- Ausgabe -> Eingabe 1:
\((d_0, d_1) \mapsto (d_0, d_1)\) für \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)
- Ausgabe -> Eingabe 2:
\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) für $\boldsymbol{d} \in [0, 2] \times [50, 79]$
Die Eingaben für die Ausgabe:
- Eingabe 1 -> Ausgabe: \((d_0, d_1) \mapsto (d_0, d_1)\) für $\boldsymbol{d} \in {\rm Dom}(input_1)$.
- Eingabe 2 -> Ausgabe: \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) für $\boldsymbol{d} \in {\rm Dom}(input_2)$.
Punkt (Ausgabe-zu-Eingabe implementiert)
Die Indexierungskarten für Punkte sind denen von Reduce ähnlich.
p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = 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 für Eingaben entspricht:
- Ausgabe -> Eingabe_1: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) für\(\boldsymbol{d} \in {\rm Dom}(output)\) und \(\boldsymbol{s} \in [0, 255]\)
- Ausgabe -> Eingabe_2: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\) für\(\boldsymbol{d} \in {\rm Dom}(output)\) und \(\boldsymbol{s} \in [0, 255]\)
Die Eingaben für die Ausgabe entsprechen:
- Eingabe_1 -> Ausgabe: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) für\(\boldsymbol{d} \in {\rm Dom}(input_1)\) und \(\boldsymbol{s} \in [0, 63]\)
- Eingabe_2 -> Ausgabe: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) für\(\boldsymbol{d} \in {\rm Dom}(input_2)\) und \(\boldsymbol{s} \in [0, 127]\)
Fenster verkleinern (TBD)
Auflage (TBD)
Maps for Fusion indexieren
Die Indexierungskarte für Fusionsvorgänge ist eine Zusammenstellung von Indexierungskarten für jeden Vorgang im Cluster. Es kann vorkommen, dass einige Eingaben mit unterschiedlichen Zugriffsmustern mehrmals gelesen werden.
Eine Eingabe, mehrere Indexierungskarten
Hier ist ein Beispiel für: \(p_0 + p_0^T\)
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 Indexierungszuordnungen von Ausgabe zu Eingabe für p0 lauten $(d_0, d_1) \mapsto (d_0, d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$. Dies bedeutet, dass wir zum Berechnen eines Elements der Ausgabe möglicherweise den Eingabeparameter zweimal lesen müssen.
Eine Eingabe, deduplizierte Indexierungskarte
Es gibt Fälle, in denen die Indexierungskarten tatsächlich identisch sind, obwohl dies nicht sofort offensichtlich 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 add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}
Die Indexierungszuordnung von Ausgabe zu Eingabe für p0 ist in diesem Fall nur $(d_0, d_1, d_2) \mapsto (d_2, d_0, d_1)$.
SoftMax
Die Indexierung von Ausgabe zu Eingabe ist für parameter 0 für Softmax zugeordnet:
- \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
- \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)
for \(\boldsymbol{d} \in {\rm Dom}(output)\) and \(\boldsymbol{s} \in [0, 124]\)bezieht sich auf die innerste Dimension der Eingabe.
Kartenvereinfacher indexieren
Mit der Standardvereinfachung für mlir::AffineMap-Upstream können keine Annahmen über die Dimensions-/Symbolbereiche getroffen werden. Daher können Ausdrücke mit mod und div nicht effizient vereinfacht werden.
Wir können das Wissen über die Unter- und Obergrenze der Unterausdrücke in den affinen Maps nutzen, um sie noch weiter zu vereinfachen.
Die Vereinfachungsfunktion kann die folgenden Ausdrücke umschreiben.
- \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) für $\boldsymbol{d} \in [0, 6] \times [0, 14]\( becomes \)(d_0, d_1) \mapsto (d_0, d_1)$
- $(d_0, d_1, d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_0, \, d_0, \, d_1, d_9, d_1, d_1\( for \)\( becomes \)
- $(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod 8)\( for \)d_i \in [0, 9]\( becomes \)(d_0, d_1, +to d_2) \
- \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) für $\boldsymbol{d} \in [0, 9] \times [0, 10]\( becomes \)(d_0, d_1) \mapsto (d_0)$.
Durch die Indexierung der Kartenvereinfachung können wir nachvollziehen, dass sich einige der verketteten Umformen 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 Zusammensetzung der Indexierungskarten und deren Vereinfachung erhalten wir
\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).