Instruções de HLO assíncrona

1. Adicionar operações assíncronas ao HLO é complicado (por exemplo, all-reduce-start e all-reduce-done).
2. As divisões de início e conclusão podem ser inadequadas para alguns dos usos assíncronos casos de uso diferentes.

Para atingir a primeira deficiência, propomos introduzir um último conjunto de novas códigos de operação assíncronas: kAsyncStart, kAsyncUpdate e kAsyncDone. A ideia é criar um código de operação assíncrono genérico que possa ajustar qualquer instrução HLO. A operação real que será executada de forma assíncrona será codificada usando uma computação chamada que tem somente a instrução como raiz e qualquer parâmetros para entradas. A manipulação de buffer de entrada/saída em trânsito e atribuição de alias podem ser compartilhados para qualquer operação assíncrona. As instruções de início assíncrono forma de saída será uma tupla dos operandos de entrada, valores de saída e quaisquer estado intermediário necessário para a async-update ou a async-done instruções.

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=”foo”
}

%async-start = (f32[64], f32[32], s32[]) async-start(f32[64] %operand),
                                         calls=%async_op
%async-done = f32[32] async-done((f32[64], f32[32], s32[]) %async-start)

Na representação acima, apenas async-start tem um cálculo chamado, já que é trivial descobrir o que async-done faz seguindo o operando para encontrar o async-start correspondente para encontrar o cálculo chamado.

Observe também que o primeiro elemento na tupla de saída dos aliases async-start com a operando para que o buffer permaneça ativo até pelo menos a instrução async-done. Da mesma forma, o segundo elemento cria um alias com a saída de async-done e o o terceiro elemento é o estado de contexto usado para acompanhar os operação assíncrona. Essa representação também aceita vários tensores em a entrada e/ou saída da operação assíncrona e a atribuição de alias funciona da mesma maneira:

%async_op {
  %param0 = f32[64] parameter(0)
  %param1 = f32[64] parameter(1)
  ROOT %op = (f32[32], f32[32]) op(f32[64] %param0, f32[64] %param1),
                                op_specific_attr=”foo”
}

%async-start = ((f32[64], f32[64]), (f32[32], f32[32]), s32[])
               async-start(f32[64] %operand0, f32[64] %operand1),
               calls=%async_op
%async-done = (f32[32], f32[32]) async-done(%async-start)

Além disso, a operação pode ser decomposta em zero ou mais async-update. etapas que realizam cálculos intermediários. A atribuição de alias de entrada/saída funciona da mesma forma com a instrução async-update, e cada async-start e As instruções de async-update precisam ter um usuário que seja outro async-update ou um async-done:

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=”foo”
}

%async-start = (f32[64], f32[32], s32[]) async-start(f32[64] %operand),
                                         calls=%async_op
%async-update0 = (f32[64], f32[32], s32[]) async-update(
                           (f32[64], f32[32], s32[]) %async-start)
%async-update1 = (f32[64], f32[32], s32[]) async-update(
                           (f32[64], f32[32], s32[]) %async-update0)
%async-done = f32[32] async-done((f32[64], f32[32], s32[]) %async-update1)

Açúcar de sintaxe

Como há um cálculo separado para definir a operação que será realizada de forma assíncrona é um pouco complicado, também propomos uma sintaxe mais leve imprime e analisa automaticamente as operações assíncronas como se fossem de primeira classe códigos de operação. A ideia é tratar os sufixos “-start”, “-update” e “-done” especialmente criando automaticamente a computação e a instrução (sem a ) ao analisar. Por exemplo, o snippet de código acima pode ter uma formatação para o seguinte, e os dois podem ser analisados na mesma representação:

%op-start = (f32[64], f32[32], s32[]) op-start(f32[64] %operand),
                                      op_specific_attr=”foo”
%op-update0 = (f32[64], f32[32], s32[]) op-update(
                        (f32[64], f32[32], s32[]) %op-start),
                        op_specific_attr=”foo”
%op-update1 = (f32[64], f32[32], s32[]) op-update(
                        (f32[64], f32[32], s32[]) %op-update0)
%op-done = f32[32] op-done((f32[64], f32[32], s32[]) %op-update1)

Para não criar ambiguidades, o verificador não permitirá que uma operação ser encapsulado com async-start se definirmos explicitamente um opcode para esse com os sufixos "-start" e/ou "-done". Esta também é uma fuga surgir caso haja alguma instrução que exija tratamento de nível HLO que não se encaixa no modelo descrito acima (por exemplo, a entrada/saída de alias reservas). Inicialmente, copy-start/copy-done collective-permute-start/collective-permute-done etc. vão continuar usando seus respectivos códigos de operação de primeira classe, em vez dos novos async-start/async-done opcodes até limparmos o código para removê-los de operações "-start"/"-done".