Cette page a été traduite par l'API Cloud Translation.

Instructions HLO asynchrones

L'ajout d'opérations asynchrones à HLO est fastidieux (par exemple, all-reduce-start et all-reduce-done).
Le fractionnement "start" et le fractionnement "terminé" peuvent ne pas être adaptés pour certaines utilisations asynchrones. cas d'utilisation.

Afin de cibler la première lacune, nous proposons d'introduire une dernière série de nouvelles codes d'opérations asynchrones: kAsyncStart, kAsyncUpdate et kAsyncDone. L'idée consiste à créer un code d'opération asynchrone générique qui peut encapsuler n'importe quelle instruction HLO. L'opération réelle qui sera effectuée de manière asynchrone est encodée à l'aide de la fonction un calcul appelé dont seule la racine est l'instruction, et toute instruction pour les entrées. Gestion des tampons d'entrée et de sortie en cours de transfert et création d'alias peuvent ensuite être partagés pour toute opération asynchrone. L'instruction de démarrage asynchrone la forme de sortie est alors un tuple d'opérandes d'entrée, de valeurs de sortie un état intermédiaire requis pour async-update ou async-done instructions.

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

Dans la représentation ci-dessus, seule async-start possède un calcul appelé puisque il est facile de trouver ce que fait async-done en suivant son opérande pour trouver la fonction async-start correspondante pour trouver le calcul appelé.

Notez également que le premier élément du tuple de sortie des alias async-start avec le paramètre opérande, de sorte que le tampon reste actif jusqu'à ce qu'il obtienne au moins l'instruction asynchrone terminée. De même, le deuxième élément est associé à la sortie async-done, et l'élément Le troisième élément est l'état de contexte utilisé pour garder la trace opération asynchrone. Cette représentation accepte également plusieurs Tensors dans l'entrée et/ou la sortie de l'opération asynchrone et la création d'alias fonctionne de la même manière. méthode:

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

De plus, l'opération peut être décomposée en zéro ou plusieurs async-update. étapes qui effectuent des calculs intermédiaires. La création d'alias en entrée/sortie fonctionne de la même manière avec l'instruction async-update, et chaque async-start et L'un des utilisateurs dans les instructions async-update doit être un autre utilisateur async-update ou 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)

Syntaxe "sucre"

Puisque le fait d'avoir un calcul distinct pour définir l'opération qui sera effectuée de manière asynchrone est un peu fastidieuse, nous proposons également un sucre syntaxique pour imprimer et analyser automatiquement les opérations asynchrones comme s'il s'agissait d'opérations de première classe des opérations. L'idée est de traiter les suffixes "-start", "-update" et "-done" notamment en créant automatiquement le calcul et l'instruction (sans ) lors de l'analyse. Par exemple, l'extrait de code ci-dessus peut être assez long à ce qui suit, et les deux peuvent être analysés dans la même représentation:

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

Afin de ne pas créer d’ambiguïté, le vérificateur n’autorise pas une opération à être encapsulé avec "async-start" si nous avons explicitement défini un code d'opération pour cette avec les suffixes "-start" et/ou "-done". C'est aussi une échappatoire au cas où des instructions nécessiteraient un traitement au niveau HLO ne correspond pas au modèle décrit ci-dessus (par exemple, les entrées/sorties de création d'alias tampons de mémoire). Au départ, copy-start/copy-done, collective-permute-start/collective-permute-done, etc. continueront d'utiliser leurs opcodes de première classe respectifs au lieu des nouveaux async-start/async-done opération(s) jusqu'à ce que nous nettoyions le code pour les supprimer Codes d'opérations "-start"/"-done".