Création d'alias dans XLA

Ce document décrit l'API de création d'alias XLA, qui vous permet de spécifier le créer un alias entre les tampons d'entrée et de sortie lors de la création d'un programme XLA.

Définir les alias au moment de la compilation

Prenons l'exemple d'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 alloue deux tampons de 4 octets: un pour l'entrée %p et un pour l'entrée pour la sortie %out.

Cependant, il est souvent souhaitable d'effectuer la mise à jour sur place (par exemple, si dans l'interface qui génère l'expression. La variable d'entrée n'est plus active. après le calcul, par exemple par incréments de p++).

Pour effectuer efficacement une telle mise à jour, vous pouvez spécifier la création d'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 indique que l'ensemble de la sortie (signalée par {}) est associée à la paramètre d'entrée 0.

Pour spécifier la création d'alias par programmation, consultez la XlaBuilder::SetUpAlias API.

Définir les alias au moment de l'exécution

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

Les tampons d'entrée du programme sont encapsulés ExecutionInput qui à son tour contiennent une arborescence de MaybeOwningDeviceMemory. Si la mémoire est spécifié comme owning (la propriété du tampon est transmise à l'environnement d'exécution XLA), le tampon est attribué et la mise à jour est exécutée sur place, demandé par l'API de création d'alias au moment de la compilation.

Toutefois, si le tampon créé au moment de la compilation n'est pas donné à la protection contre la copie est activée: un tampon de sortie supplémentaire O est alloué, et le contenu du tampon d'entrée P qui devait être créé avec un alias sont copiés dans O (de sorte que le programme peut s'exécuter comme si le tampon O était donné au moment de l'exécution).