En este documento, se describe la API de alias de XLA, que te permite especificar la la asignación de alias entre los búferes de entrada y salida cuando se compila un programa de XLA.
Define alias en el tiempo de compilación
Por ejemplo, considera 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 %p de entrada y otro para la entrada.
para la salida %out.
Sin embargo, suele ser conveniente realizar la actualización in situ (por ejemplo, si
En el frontend que genera la expresión, la variable de entrada ya no está activa.
después del cálculo, como en el incremento p++).
Para realizar esta actualización de forma eficiente, puedes 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 por {}) tiene un alias al
parámetro de entrada 0.
Para especificar la asignación de alias de forma programática, consulta la
XlaBuilder::SetUpAlias
en la API de Cloud.
Definición de alias en el entorno de ejecución
El alias definido en el paso anterior se especifica durante la compilación.
Durante la ejecución, puedes usar el
LocalClient::RunAsync
API para elegir si donar el «colchón».
Los búferes de entrada al programa se unen
ExecutionInput,
que, a su vez, contienen un árbol de MaybeOwningDeviceMemory. Si la memoria
especificado como owning (la propiedad del búfer se pasa al tiempo de ejecución de XLA)
se dona el búfer y se ejecuta la actualización en el lugar,
que solicita la API de alias en el tiempo de compilación.
Sin embargo, si el búfer con alias en el tiempo de compilación no se dona en
durante el tiempo de ejecución, se activa la protección contra copias: se asigna un búfer de salida adicional O,
y se copia el contenido del búfer de entrada P que debía tener un alias.
a O (de manera que el programa pueda ejecutarse
como si el búfer O estuviera
se donan durante el tiempo de ejecución).