Petunjuk HLO Asinkron

  1. Menambahkan operasi asinkron ke HLO rumit (yaitu all-reduce-start dan all-reduce-done).
  2. Pemisahan yang dimulai dan selesai mungkin tidak memadai untuk beberapa penggunaan asinkron penggunaan.

Untuk menargetkan kekurangan pertama, kami mengusulkan untuk memperkenalkan satu set opcode asinkron: kAsyncStart, kAsyncUpdate, dan kAsyncDone. Ide adalah membuat opcode asinkron generik yang dapat menggabungkan instruksi HLO. Operasi sebenarnya yang akan dijalankan secara asinkron akan dienkode menggunakan suatu komputasi yang disebut yang hanya memiliki instruksi sebagai {i>root<i}-nya dan setiap parameter untuk input. Penanganan dan aliasing buffer input/output dalam proses kemudian dapat dibagikan untuk setiap operasi asinkron. Petunjuk async-start bentuk output, lalu akan menjadi tuple dari operand input, nilai output, dan status perantara yang diperlukan untuk async-update atau async-done petunjuk.

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

Dalam representasi di atas, hanya async-start yang memiliki komputasi yang disebut karena tidaklah sulit menemukan apa yang dilakukan async-done dengan mengikuti operand-nya ke menemukan async-start yang sesuai untuk menemukan komputasi yang disebut.

Perhatikan juga bahwa elemen pertama dalam tuple output dari alias async-start dengan operand, sehingga buffer tetap aktif hingga setidaknya instruksi asinkron selesai. Demikian pula, alias elemen kedua dengan output async-done, dan elemen ketiga adalah status konteks yang digunakan untuk melacak dan asinkron. Representasi ini juga mendukung beberapa tensor di input dan/atau output operasi asinkron dan {i>aliasing<i} bekerja dengan cara yang sama cara:

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

Selain itu, op dapat diurai lebih lanjut menjadi nol atau beberapa async-update langkah-langkah yang melakukan komputasi perantara. Aliasing input/output berfungsi dengan cara yang sama dengan instruksi async-update, serta setiap async-start dan Petunjuk async-update harus memiliki satu pengguna yang merupakan pengguna lain async-update atau 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)

sugar sintaksis

Karena memiliki komputasi terpisah untuk mendefinisikan operasi yang akan dilakukan secara asinkron sedikit rumit, kami juga mengusulkan sugar {i>syntax<i} untuk secara otomatis mencetak dan mengurai operasi asinkron seolah-olah operasi tersebut kelas satu opcode. Tujuannya adalah untuk menangani akhiran “-start”, “-update”, dan “-done” khususnya dengan secara otomatis membuat komputasi dan instruksi (tanpa akhiran) saat mengurai. Misalnya, cuplikan kode di atas bisa sangat dicetak menjadi representasi berikut dan keduanya dapat diurai menjadi representasi yang sama:

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

Agar tidak menimbulkan ambiguitas, pemverifikasi tidak akan mengizinkan operasi untuk digabungkan dengan async-start jika kita secara eksplisit menentukan opcode untuk operasi dengan akhiran “-start” dan/atau “-done”. Ini juga merupakan escape menetas jika kami memiliki petunjuk yang memerlukan perawatan tingkat HLO yang tidak sesuai dengan model yang dijelaskan di atas (misalnya input/output aliasing {i>buffer<i}). Jadi, awalnya, copy-start/copy-done, collective-permute-start/collective-permute-done dll. akan terus menggunakan opcode kelas satu mereka masing-masing alih-alih opcode async-start/async-done hingga kita membersihkan kode untuk menghapusnya opcode “-start”/”-done”.