Ce document décrit l'API de création d'alias XLA, qui vous permet de spécifier la création d'alias entre les tampons d'entrée et de sortie lors de la création d'un programme XLA.
Définir la création d'alias au moment de la compilation
Prenons l'exemple d'un module HLO simple 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 l'autre 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, comme lors de l'incrément 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 spécifie que l'intégralité de la sortie (marquée par {}) est associée au paramètre d'entrée 0.
Pour spécifier la création d'alias de manière programmatique, consultez l'API XlaBuilder::SetUpAlias.
Définir la création d'alias au moment de l'exécution
La création d'alias définie à l'étape précédente est spécifiée lors de la compilation.
Lors de l'exécution, vous pouvez utiliser l'API LocalClient::RunAsync pour choisir de faire don du tampon.
Les tampons d'entrée du programme sont encapsulés dans des ExecutionInput, qui à leur tour contiennent une arborescence de MaybeOwningDeviceMemory. Si la mémoire est spécifiée comme propriétaire (la propriété du tampon est transmise à l'environnement d'exécution XLA), le tampon est effectivement partagé et la mise à jour est exécutée sur place, comme demandé par l'API de création d'alias au moment de la compilation.
Toutefois, si le tampon associé au moment de la compilation n'est pas donné au moment de l'exécution, la fonctionnalité copy-protection s'active: un tampon de sortie supplémentaire O est alloué et le contenu du tampon d'entrée P devant être associé à un alias est copié dans O (le programme peut donc s'exécuter comme si le tampon O avait été donné au moment de l'exécution).