Este documento describe la API de aliasing para XLA: al crear un programa XLA, puede especificar el aliasing deseado entre los búferes de entrada y salida.
Definición de alias en tiempo de compilación
Por ejemplo, considere un módulo HLO trivial que simplemente agrega 1 a su entrada:
HloModule increment
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
Este módulo asignará dos búferes de 4 bytes: uno para la entrada %p otro para la salida %out .
Sin embargo, a menudo es deseable realizar la actualización en el lugar (por ejemplo, si en la interfaz que genera la expresión, la variable de entrada ya no está viva después del cálculo, como en el incremento p++ ).
Para realizar una actualización de este tipo de manera eficiente, puede especificar el alias de entrada:
HloModule increment, input_output_alias={ {}: 0 }
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
El formato especifica que toda la salida (marcada con {} ) tiene como alias el parámetro de entrada 0 .
Consulte la API XlaBuilder::SetUpAlias para especificar el alias mediante programación.
Definición de alias en tiempo de ejecución
El alias definido en el paso anterior se especifica durante la compilación . Durante la ejecución, puede elegir si donar el búfer utilizando la API LocalClient::RunAsync .
Los búferes de entrada al programa están empaquetados en ExecutionInput , que a su vez contiene un árbol de MaybeOwningDeviceMemory . Si la memoria se especifica como propietaria (la propiedad del búfer se pasa al tiempo de ejecución de XLA), el búfer en realidad se dona y la actualización se ejecuta en el lugar, según lo solicite la API de alias en tiempo de compilación.
Sin embargo, si el búfer que tiene un alias en tiempo de compilación no se dona en tiempo de ejecución, se activa la protección contra copia : se asigna un búfer de salida adicional O , y el contenido del búfer de entrada P que estaba destinado a ser alias se copia en O (de manera que el programa puede ejecutarse efectivamente como si el búfer O fuera donado en tiempo de ejecución).
Interoperabilidad frontend
TF / XLA
En los clústeres del programa TensorFlow compilados con XLA, todas las actualizaciones de las variables de recursos tienen un alias en el tiempo de compilación (el alias en el tiempo de ejecución depende de si algo más tiene una referencia al tensor de la variable del recurso).