Esta página foi traduzida pela API Cloud Translation.
Switch to English

Aliasing no XLA

Este documento descreve a API de aliasing para XLA: ao construir um programa XLA, você pode especificar o aliasing desejado entre os buffers de entrada e saída.

Definição de apelido em tempo de compilação

Por exemplo, considere um módulo HLO trivial que simplesmente adiciona 1 à sua entrada:

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Este módulo alocará dois buffers de 4 bytes: um para a entrada %p , e um para a saída %out .

No entanto, muitas vezes é desejável realizar a atualização no local (por exemplo, se no frontend que gera a expressão a variável de entrada não estiver mais viva após o cálculo, como no incremento p++ ).

Para realizar essa atualização com eficiência, você pode especificar o 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 {} ) é um alias para o parâmetro de entrada 0 .

Consulte a API XlaBuilder::SetUpAlias para especificar o aliasing de forma programática.

Definindo aliasing em tempo de execução

O aliasing definido na etapa anterior é especificado durante a compilação . Durante a execução, você pode escolher se realmente quer doar o buffer usando a API LocalClient::RunAsync .

Os buffers de entrada para o programa são agrupados em ExecutionInput , que por sua vez contém uma árvore de MaybeOwningDeviceMemory . Se a memória for especificada como proprietária (a propriedade do buffer é passada para o tempo de execução do XLA), o buffer é realmente doado e a atualização é executada no local, conforme solicitado pela API de aliasing em tempo de compilação.

Se, no entanto, o buffer que recebe um alias em tempo de compilação não é doado em tempo de execução, a proteção contra cópia é ativada: um buffer de saída extra O é alocado e o conteúdo do buffer de entrada P que deveria ter um alias é copiado para O (de forma eficaz, o programa pode ser executado como se o buffer O tivesse sido doado em tempo de execução).

Interoperabilidade de front-end

TF / XLA

Em clusters do programa TensorFlow compilado com XLA, todas as atualizações de variáveis ​​de recursos recebem um alias no tempo de compilação (o alias no tempo de execução depende se alguma outra coisa contém uma referência ao tensor da variável de recurso).