En este documento, se describe la API de alias de XLA, que te permite especificar el alias entre los búferes de entrada y salida cuando compilas un programa de XLA.
Cómo definir 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 %out de salida.
Sin embargo, a menudo es 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 una actualización de este tipo de manera 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 para el parámetro de entrada 0.
Para especificar la asignación de alias de manera programática, consulta la API de XlaBuilder::SetUpAlias.
Cómo definir 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 la API de LocalClient::RunAsync para elegir si donar el búfer.
Los búferes de entrada del programa se unen en ExecutionInput, que, a su vez, contienen un árbol de MaybeOwningDeviceMemory. Si la memoria se especifica como propiedad (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, según lo solicite la API de alias en tiempo de compilación.
Sin embargo, si el búfer al que se le asigna un alias en el tiempo de compilación no se dona en el tiempo de ejecución, se activa copy-protection: se asigna un búfer de salida adicional O y el contenido del búfer de entrada P al que se debía asignar un alias se copia en O (de manera efectiva, el programa puede ejecutarse como si el O del búfer se donara en el tiempo de ejecución).