Este documento descreve a API de alias do XLA, que permite especificar atribuir alias entre os buffers de entrada e de saída ao criar um programa XLA.
Como definir a atribuição de alias no momento da compilação
Por exemplo, considere um módulo HLO trivial que simplesmente adiciona 1 à entrada:
HloModule increment
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
Este módulo vai alocar dois buffers de 4 bytes: um para a entrada %p e outro
para a saída %out.
No entanto, muitas vezes é desejável realizar a atualização no local (por exemplo, se
no front-end que gera a expressão, a variável de entrada não está mais ativa
após o cálculo, como no incremento p++).
Para realizar essa atualização de forma eficiente, especifique a atribuição de 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)
}
O formato especifica que toda a saída (marcada por {}) recebe um alias para a
parâmetro de entrada 0.
Para especificar a atribuição de alias de forma programática, consulte a
XlaBuilder::SetUpAlias
API.
Como definir a atribuição de alias no ambiente de execução
A atribuição de alias definida na etapa anterior é especificada durante a compilação.
Durante a execução, você pode usar o método
LocalClient::RunAsync
API para escolher se quer doar o buffer.
Os buffers de entrada do programa estão agrupados
ExecutionInputs,
que, por sua vez, contêm uma árvore de MaybeOwningDeviceMemory. Se a memória for
especificado como proprietário (a propriedade do buffer é passada para o ambiente de execução do XLA);
o buffer é realmente doado e a atualização é executada no local,
solicitado pela API de atribuição de alias no tempo de compilação.
No entanto, se o buffer com alias no tempo de compilação não for doado em
ambiente de execução, copy-protection é iniciado: um buffer de saída extra O é alocado,
e o conteúdo do buffer de entrada P que deveria ter um alias é copiado
em O, para que o programa possa ser executado de forma efetiva como se o buffer O estivesse
doados no tempo de execução).