Atribuição de alias no XLA

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 ExecutionInputs, 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.