Instrucciones de HLO asíncrono

1. Agregar operaciones asíncronas a HLO es un proceso engorroso (es decir, all-reduce-start y all-reduce-done).
2. La división de inicio y finalización puede no ser adecuada para algunos de los usos asíncronos. diferentes.

Para abordar la primera deficiencia, proponemos introducir un último conjunto de nuevos Códigos de operación asíncronos: kAsyncStart, kAsyncUpdate y kAsyncDone La idea es crear un código de operación asíncrono y genérico que pueda unir cualquier instrucción HLO. La operación real que se realizará de forma asíncrona se codificará con un cálculo llamado que solo tiene la instrucción como su raíz y cualquier parámetros para las entradas. El alias y el manejo del búfer de entrada y salida en tránsito se pueden compartir para cualquier operación asíncrona. La instrucción de inicio asíncrono forma de salida será una tupla de operandos de entrada, valores de salida y cualquier el estado intermedio que se necesita para los objetos async-update o async-done instrucciones.

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

En la representación anterior, solo async-start tiene un cálculo llamado, ya que es trivial encontrar lo que hace async-done siguiendo su operando a busca el async-start correspondiente para encontrar el procesamiento llamado.

También ten en cuenta que el primer elemento de la tupla de salida de alias async-start con el operando, de modo que el búfer permanece activo al menos hasta la instrucción async-done. De manera similar, el segundo elemento se alias con el resultado de async-done y el El tercer elemento es el estado del contexto que se usa para hacer un seguimiento de la una operación asíncrona. Esta representación también admite varios tensores en la entrada o salida de la operación asíncrona y la asignación de alias funcionan de la misma forma:

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

Además, la op se puede descomponer en cero o más async-update. pasos que realizan cálculos intermedios. La asignación de alias de entrada y salida funciona de la misma manera con la instrucción async-update y cada async-start y Las instrucciones async-update deben tener un usuario distinto async-update o 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)

Azúcar de sintaxis

Dado que tener un cálculo independiente para definir la operación que se de forma asíncrona es un poco engorrosa, también proponemos una sintaxis edulcorada imprimir y analizar automáticamente operaciones asíncronas como si fueran de primera clase códigos de operación. La idea es tratar los sufijos “-start”, “-update” y “-done”. en especial mediante la creación automática del procesamiento y la instrucción (sin el sufijo) durante el análisis. Por ejemplo, el fragmento de código anterior puede imprimirse con formato estilístico. como se muestra a continuación, y ambas se pueden analizar en la misma representación:

%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 no crear ambigüedades, el verificador no permitirá que una operación se unen con async-start si definimos explícitamente un código de operación para ese con los sufijos “-start” o “-done”. Esto también es un escape en caso de que tengamos instrucciones que requieran un tratamiento a nivel de HLO no se ajusta al modelo descrito anteriormente (p.ej., la entrada/salida de alias tiempos de reserva). Por lo tanto, inicialmente, copy-start/copy-done, collective-permute-start, collective-permute-done, etc., seguirán usando sus respectivos códigos de operación de primera clase, en lugar del nuevo Códigos de operación async-start/async-done hasta que limpiemos el código para quitarlos. códigos de operación “-start”/”-done”.