Anleitung für Async HLO

1. Das Hinzufügen von asynchronen Vorgängen zu HLO ist umständlich (d.h. all-reduce-start und all-reduce-done).
2. 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.