Atribuição de alias no XLA

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).