Instrukcje asynchronicznego HLO

1. Dodawanie operacji asynchronicznych do HLO jest uciążliwe (np. all-reduce-start i all-reduce-done).
2. Podział startu i zakończenia może być nieodpowiedni w przypadku części zastosowań asynchronicznych przypadków.

Aby dotrzeć do pierwszego niedociągnięcia, proponujemy wprowadzenie ostatniego zestawu nowych asynchroniczne kody operacji: kAsyncStart, kAsyncUpdate i kAsyncDone. Pomysł jest utworzenie ogólnego asynchronicznego kodu operacji, który opakowuje dowolną instrukcję HLO. Rzeczywista operacja, która zostanie wykonana asynchronicznie, zostanie zakodowana za pomocą czyli tzw. obliczenia, w których pierwiastkiem jest tylko instrukcja dla danych wejściowych. Obsługa bufora wejścia/wyjścia i stosowanie aliasów który można udostępnić dla dowolnej operacji asynchronicznej. Instrukcje uruchamiania asynchronicznego kształt wyjściowy będzie krotką operandów wejściowych, wartości wyjściowych i dowolnych stan pośredni wymagany przez async-update lub async-done za instrukcje.

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

W powyższej reprezentacji tylko async-start wywołało funkcję obliczeń od łatwo jest znaleźć działanie funkcji async-done, podążając za jej argumentem znajdź odpowiednią wartość async-start, aby znaleźć tzw. obliczenia.

Pamiętaj też o tym że pierwszy element w krotce wyjściowej async-start zawiera alias z parametrem operand, dzięki czemu bufor pozostanie aktywny do momentu wykonania instrukcji (co najmniej do czasu wykonania instrukcji asynchronicznej). Podobnie drugi element aliasuje dane wyjściowe async-done, a parametry trzeci element to stan kontekstu służący do śledzenia operacji asynchronicznej. Ta reprezentacja obsługuje również wiele tensorów w dane wejściowe lub wyjściowe operacji asynchronicznej i aliasy działają tak samo sposób:

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

Dodatkowo działanie można podzielić na zero lub więcej async-update. do wykonywania obliczeń pośrednich. Aliasy wejściowe/wyjściowe działają w ten sam sposób z instrukcją async-update i każdym elementem async-start i Instrukcje typu async-update muszą należeć do jednego użytkownika async-update lub 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)

Składnia cukru

Ponieważ istnieją oddzielne obliczenia w celu zdefiniowania operacji, która zostanie asynchronicznie jest uciążliwe, ale oferujemy też cukier składniowy, automatycznie drukuje i analizuje operacje asynchroniczne tak, jakby były to operacje pierwszej klasy opcodes. Chodzi o traktowanie sufiksów „-start”, „-update” i „-done” specjalnie przez automatyczne tworzenie obliczeń i instrukcji (bez sufiksu). Na przykład ten fragment kodu może być na następujące. Oba te współczynniki można interpretować jako takie same:

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

Aby uniknąć niejasności, weryfikator nie zezwoli na operację powinna być zapakowana z parametrem async-start, jeśli jednoznacznie zdefiniowaliśmy kod operacji z sufiksami „-start” i/lub „-done”. To też może być ucieczka w razie potrzeby w razie potrzeby nie mieści się w modelu opisanym powyżej (np. aliasy wejścia/wyjścia bufory). Więc początkowo copy-start/copy-done, collective-permute-start/collective-permute-done itp. będzie nadal używać odpowiednie kody opcode pierwszej klasy zamiast nowego Kody operacji async-start/async-done, dopóki nie pozbądźmy się kodu, aby je usunąć Kody operacji „-start”/”-done”.