Este documento descreve a API de alias do XLA, que permite especificar a alias entre os buffers de entrada e saída ao criar um programa do XLA.
Como definir a atribuição de alias durante a 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)
}
Esse módulo 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 estiver mais ativa
após o cálculo, como no incremento p++
.
Para executar essa atualização de forma eficiente, é possível especificar 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 com {}
) tem o alias definido como o
parâmetro de entrada 0
.
Para especificar a atribuição de alias de forma programática, consulte a
API
XlaBuilder::SetUpAlias
.
Como definir a atribuição de alias no tempo de execução
A atribuição de alias definida na etapa anterior é especificada durante a compilação.
Durante a execução, é possível usar a API
LocalClient::RunAsync
para escolher se quer doar o buffer.
Os buffers de entrada para o programa são encapsulados em
ExecutionInput
s,
que, por sua vez, contêm uma árvore de MaybeOwningDeviceMemory
. Se a memória for
especificada como propriedade (a propriedade do buffer é transmitida ao ambiente de execução do XLA),
o buffer será realmente doado e a atualização será executada no local, conforme
solicitado pela API de atribuição de alias no tempo de compilação.
Entretanto, se o buffer com alias no tempo de compilação não for doado no
tempo de execução, a copy-protection será iniciada: um buffer de saída extra O
será alocado
e o conteúdo do buffer de entrada P
, que deveria ter o alias, será copiado
em O
. Dessa forma, o programa poderá ser executado como se o buffer O
tivesse sido
doado no momento da execução.