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

Aliasing dans XLA

Ce document décrit l'API d'aliasing pour XLA: lors de la construction d'un programme XLA, vous pouvez spécifier l'aliasing souhaité entre les tampons d'entrée et de sortie.

Définition de l'alias au moment de la compilation

Par exemple, considérons un module HLO trivial qui ajoute simplement 1 à son entrée:

 HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}
 

Ce module allouera deux tampons de 4 octets: un pour l'entrée %p et un pour la sortie %out .

Cependant, il est souvent souhaitable d'effectuer la mise à jour sur place (par exemple, si dans le frontend générant l'expression la variable d'entrée n'est plus vivante après le calcul, comme dans l'incrément p++ ).

Pour effectuer une telle mise à jour efficacement, vous pouvez spécifier l'alias d'entrée:

 HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}
 

Le format spécifie que la sortie entière (marquée par {} ) est aliasée sur le paramètre d'entrée 0 .

Consultez l'API XlaBuilder::SetUpAlias pour spécifier l'alias par programmation.

Définition d'alias au moment de l'exécution

L'alias défini à l'étape précédente est spécifié lors de la compilation . Pendant l'exécution, vous pouvez choisir de donner ou non le tampon à l'aide de l'API LocalClient::RunAsync .

Les tampons d'entrée du programme sont encapsulés dans ExecutionInput , qui à son tour contient une arborescence de MaybeOwningDeviceMemory . Si la mémoire est spécifiée comme propriétaire (la propriété du tampon est transmise au runtime XLA), le tampon est en fait donné et la mise à jour est exécutée sur place, comme demandé par l'API d'alias de compilation.

Si, cependant, le tampon qui est aliasé au moment de la compilation n'est pas donné au moment de l'exécution, la protection contre la copie intervient: un tampon de sortie supplémentaire O est alloué, et le contenu du tampon d'entrée P qui était censé être aliasé est copié dans O (si efficacement le programme peut s'exécuter comme si le tampon O avait été donné au moment de l'exécution).

Interopérabilité frontale

TF / XLA

Dans les clusters du programme TensorFlow compilé avec XLA, toutes les mises à jour de variable de ressource sont aliasées au moment de la compilation (l'alias au moment de l'exécution dépend de si autre chose contient une référence au tenseur de variable de ressource).