- Das Hinzufügen von asynchronen Vorgängen zu HLO ist umständlich (d.h.
all-reduce-start
undall-reduce-done
). - Die Aufteilung von "Start" und "Done" ist für einen Teil der asynchronen Verwendung möglicherweise nicht ausreichend. Cases.
Um die erste Schwachstelle abzubauen, schlagen wir vor, eine letzte Reihe neuer
Asynchrone Opcodes: kAsyncStart
, kAsyncUpdate
und kAsyncDone
. Die Idee
besteht darin, einen generischen asynchronen Opcode zu erstellen, der jede HLO-Anweisung umschließen kann.
Der eigentliche Vorgang, der asynchron durchgeführt wird, wird mit
eine namens Berechnung, die nur die Anweisung als Stamm hat und
Parameter für Eingaben. Verarbeitung und Aliasing des In-Flight-Eingabe-/-Ausgabepuffers
kann dann für beliebige asynchrone Vorgänge freigegeben werden. Die Anweisung
Die Ausgabeform ist dann ein Tupel der Eingabeoperanden, Ausgabewerte und beliebigen
Zwischenstatus, der für async-update
oder async-done
erforderlich ist
Anleitung.
%async_op {
%param0 = f32[64] parameter(0)
ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=”foo”
}
%async-start = (f32[64], f32[32], s32[]) async-start(f32[64] %operand),
calls=%async_op
%async-done = f32[32] async-done((f32[64], f32[32], s32[]) %async-start)
In der obigen Darstellung verfügt nur async-start
über eine aufgerufene Berechnung, seit
ist es einfach, herauszufinden, was async-done
tut, indem man seinem Operanden folgt
den entsprechenden async-start
finden, um die aufgerufene Berechnung zu finden.
Weitere Hinweise
dass das erste Element im Ausgabe-Tupel von async-start
-Aliassen mit dem
damit der Puffer bis mindestens zur Anweisung "async-done" aktiv bleibt.
Ähnlich verhält es sich mit dem Alias des zweiten Elements mit der Ausgabe von async-done
.
Das dritte Element ist der Kontextstatus, mit dem die
asynchronem Vorgang. Diese Darstellung unterstützt auch mehrere Tensoren in
Die Eingabe und/oder Ausgabe des asynchronen Vorgangs und das Aliasing funktionieren gleich
Weg:
%async_op {
%param0 = f32[64] parameter(0)
%param1 = f32[64] parameter(1)
ROOT %op = (f32[32], f32[32]) op(f32[64] %param0, f32[64] %param1),
op_specific_attr=”foo”
}
%async-start = ((f32[64], f32[64]), (f32[32], f32[32]), s32[])
async-start(f32[64] %operand0, f32[64] %operand1),
calls=%async_op
%async-done = (f32[32], f32[32]) async-done(%async-start)
Darüber hinaus kann der Vorgang in null oder mehr async-update
aufgeteilt werden
Schritte, die Zwischenberechnungen durchführen. Das Eingabe-/Ausgabealiasing funktioniert,
mit der Anweisung async-update
und jeder async-start
und
async-update
-Anweisungen müssen einen Nutzer enthalten, der entweder ein anderer ist
async-update
oder ein async-done
:
%async_op {
%param0 = f32[64] parameter(0)
ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=”foo”
}
%async-start = (f32[64], f32[32], s32[]) async-start(f32[64] %operand),
calls=%async_op
%async-update0 = (f32[64], f32[32], s32[]) async-update(
(f32[64], f32[32], s32[]) %async-start)
%async-update1 = (f32[64], f32[32], s32[]) async-update(
(f32[64], f32[32], s32[]) %async-update0)
%async-done = f32[32] async-done((f32[64], f32[32], s32[]) %async-update1)
Syntax Zucker
Da eine separate Berechnung zur Definition des Vorgangs erforderlich ist, asynchron ausgeführt wird, ist ein wenig umständlich. Wir schlagen außerdem Asynchrone Vorgänge automatisch drucken und parsen, als wären sie erstklassig opcodes. Dabei sollen die Suffixe „-start“, „-update“ und „-done“ indem die Berechnung und Anweisung automatisch erstellt wird (ohne Suffix) beim Parsen. Das obige Code-Snippet kann beispielsweise elegant gedruckt werden. in das Folgende und können beide in dieselbe Darstellung geparst werden:
%op-start = (f32[64], f32[32], s32[]) op-start(f32[64] %operand),
op_specific_attr=”foo”
%op-update0 = (f32[64], f32[32], s32[]) op-update(
(f32[64], f32[32], s32[]) %op-start),
op_specific_attr=”foo”
%op-update1 = (f32[64], f32[32], s32[]) op-update(
(f32[64], f32[32], s32[]) %op-update0)
%op-done = f32[32] op-done((f32[64], f32[32], s32[]) %op-update1)
Um Unklarheiten zu vermeiden, lässt die Prüfung keinen Vorgang zu, bei dem
mit async-start umschlossen, wenn wir explizit einen Opcode dafür definiert haben.
mit dem Suffix „-start“ und/oder „-done“. Das ist auch ein Fluchtweg.
falls wir Anweisungen haben, die eine Behandlung auf HLO-Ebene erfordern,
nicht in das oben beschriebene Modell passt (z.B. die Aliasing-Eingabe/-Ausgabe)
Puffer). Zunächst copy-start
/copy-done
,
collective-permute-start
/collective-permute-done
usw. wird weiterhin verwendet
ihre jeweiligen First-Class-Opcodes anstelle des neuen
async-start
/async-done
-Opcodes, bis wir den Code bereinigen und diese
„-start“/„-done“-Opcodes.