دستورالعمل های HLO Async

1. افزودن عملیات ناهمگام به HLO کار دشواری است (یعنی all-reduce-start و all-reduce-done ).
2. تقسیم‌بندی شروع و پایان ممکن است برای برخی از موارد استفاده ناهمزمان ناکافی باشد.

برای رفع اولین نقص، پیشنهاد می‌کنیم یک مجموعه نهایی از کدهای عملیاتی ناهمزمان جدید را معرفی کنیم: kAsyncStart ، kAsyncUpdate و kAsyncDone . ایده این است که یک کد عملیاتی ناهمزمان عمومی ایجاد شود که بتواند هر دستورالعمل HLO را در بر بگیرد. عملیات واقعی که به صورت ناهمزمان انجام می‌شود، با استفاده از یک محاسبه فراخوانی شده که فقط دستورالعمل را به عنوان ریشه و هر پارامتری برای ورودی‌ها دارد، کدگذاری می‌شود. سپس می‌توان مدیریت و نام‌گذاری بافر ورودی/خروجی در حین اجرا را برای هر عملیات ناهمزمان به اشتراک گذاشت. شکل خروجی دستورالعمل async-start سپس یک چندتایی از عملوندهای ورودی، مقادیر خروجی و هر حالت میانی مورد نیاز برای دستورالعمل‌های async-update یا 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-done = f32[32] async-done(((f32[64]), f32[32], s32[]) %async-start)

در نمایش بالا، فقط async-start یک محاسبه فراخوانی شده دارد، زیرا پیدا کردن کاری که async-done با دنبال کردن عملوند آن برای یافتن async-start مربوطه برای یافتن محاسبه فراخوانی شده انجام می‌دهد، امری بدیهی است.

همچنین توجه داشته باشید که اولین عنصر در تاپل خروجی async-start ، تاپل حاوی عملوندها است. عناصر این تاپل عملوند با عملوندهای مربوطه نام مستعار دارند، بنابراین بافرهای آنها حداقل تا زمان دستورالعمل async-done زنده می‌مانند. به طور مشابه، عنصر دوم با خروجی async-done نام مستعار دارد و عنصر سوم حالت زمینه‌ای است که برای پیگیری عملیات ناهمزمان استفاده می‌شود. این نمایش به طور طبیعی از چندین تنسور در ورودی و/یا خروجی عملیات ناهمزمان پشتیبانی می‌کند:

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

علاوه بر این، عملیات می‌تواند به صفر یا چند مرحله‌ی async-update که محاسبات میانی را انجام می‌دهند، تجزیه شود. نامگذاری ورودی/خروجی به همان روش با دستورالعمل async-update عمل می‌کند و هر دستورالعمل async-start و async-update باید یک کاربر داشته باشد که یا یک async-update دیگر یا یک 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)

قند نحوی

تجزیه‌گر HLO از syntax sugar پشتیبانی می‌کند تا عملیات ناهمزمان را به طور خودکار تجزیه و چاپ کند، گویی که کدهای عملیاتی درجه یک هستند. این تجزیه‌گر با پسوندهای -start ، -update و -done به طور خاص رفتار می‌کند و به طور خودکار محاسبات ناهمزمان و دستورالعمل پیچیده شده (بدون پسوند) را ایجاد می‌کند.

برای مثال، یک custom-call ناهمزمان می‌تواند به صورت زیر نوشته شود:

%cc-start = ((f32[64]), f32[32], s32[]) custom-call-start(%operand),
                                        custom_call_target="foo"
%cc-update = ((f32[64]), f32[32], s32[]) custom-call-update(%cc-start)
%result = f32[32] custom-call-done(%cc-update)

تجزیه‌کننده این را به HLO معادل زیر تبدیل می‌کند:

%async_computation {
  %p0 = f32[64] parameter(0)
  ROOT %custom-call = f32[32] custom-call(%p0), custom_call_target="foo"
}

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

این desugaring برای اکثر کدهای عملیاتی HLO (مثلاً custom-call ، dot ، all-reduce و غیره) پشتیبانی می‌شود.

استثنائات

برای جلوگیری از ایجاد ابهام، تجزیه‌گر، عملیاتی را که دارای کدهای عملیاتی صریح درجه یک هستند و با پسوندهای -start و/یا -done تعریف شده‌اند (مثلاً copy-start / copy-done ، collective-permute-start / collective-permute-done . این عملیات‌ها همچنان از کدهای عملیاتی درجه یک مربوط به خود استفاده خواهند کرد.

صحافی دیرهنگام

در برخی موارد، عملوندها (ورودی‌ها) یا خروجی‌های یک عملیات ناهمزمان، هنگام شروع عملیات، همگی در دسترس یا تخصیص داده نمی‌شوند. XLA از اتصال دیرهنگام پشتیبانی می‌کند، که به عملوندها اجازه می‌دهد تا در طول مراحل async-update به صورت افزایشی متصل شوند و خروجی‌ها نیز در طول مراحل async-update یا async-done متصل شوند.

نمایندگی در HLO

برای یک محاسبه‌ی فراخوانی شده که انتظار دارد \(N\) پارامترها، می‌توانیم اجرای ناهمزمان را با تعداد کمتری از \(N\) عملوندها. عملوندهای باقیمانده در دستورالعمل‌های بعدیِ async-update ارسال می‌شوند.

• async-start اولین را متصل می‌کند \(K\) عملوندها (\(K < N\)).
• دستورالعمل‌های async-update باقیمانده را متصل می‌کنند. \(N - K\) عملوندها

اتصال عملوندها باید از چپ به راست انجام شود. یعنی اگر یک محاسبه انتظار پارامترها را داشته باشد، \(P_0, P_1, \dots, P_{N-1}\)، آنها باید به آن ترتیب در سراسر زنجیره ناهمزمان متصل شوند.

شکل‌های async-start و async-update پارامترهای مقید شده به صورت افزایشی را منعکس می‌کنند. به طور خاص، اولین عنصر شکل چندتایی (اشکال عملوند) با مقید شدن عملوندهای بیشتر، رشد می‌کند.

اتصال خروجی مستقل از اتصال عملوند است و می‌تواند در هر مرحله‌ای از زنجیره async (چه در async-update و چه در async-done نهایی) اتفاق بیفتد.

مثال با kCall

یک محاسبه به نام %foo را در نظر بگیرید که دو پارامتر می‌گیرد:

%foo {
  %p0 = f32[] parameter(0)
  %p1 = f32[] parameter(1)
  ROOT %add = f32[] add(%p0, %p1)
}

می‌توانیم این محاسبه را به صورت غیرهمزمان فراخوانی کنیم، به این صورت که %p0 را در شروع و %p1 را در به‌روزرسانی مقید کنیم:

%call-start = ((f32[]), (), s32[]) call-start(%operand0), to_apply=%foo
%call-update = ((f32[], f32[]), f32[], s32[]) call-update(%call-start, %operand1)
%result = f32[] call-done(%call-update)

تجزیه‌کننده این را به HLO زیر تبدیل می‌کند:

%async-start = ((f32[]), (), s32[]) async-start(%operand0), calls=%foo
%async-update = ((f32[], f32[]), f32[], s32[]) async-update(%async-start, %operand1)
%result = f32[] async-done(%async-update)

خروجی‌های با محدودیت دیرهنگام

علاوه بر عملوندها (ورودی‌ها)، خروجی‌های یک عملیات ناهمزمان نیز می‌توانند با تأخیر مقید شوند. این زمانی مفید است که بافرهای خروجی در شروع عملیات شناخته شده یا تخصیص داده نشده باشند.

برای نمایش خروجی‌های با محدودیت دیرهنگام:

1. دستورالعمل async-start (یا call-start ) با یک تاپل () خالی در اندیس ۱ شکل خروجی خود (شکاف نتیجه) تعریف می‌شود.
2. یک دستورالعمل async-update (یا call-update ) بعدی، شکل خروجی واقعی را در اندیس ۱ مشخص می‌کند و جایگزین تاپل خالی می‌شود.
3. به طور جایگزین، خروجی می‌تواند در انتهای زنجیره توسط دستورالعمل async-done (یا call-done ) مقید شود، که شکل خروجی نهایی را برمی‌گرداند. این کار می‌تواند صرف نظر از وجود مراحل میانی async-update در زنجیره انجام شود.

مثال با async-update

// Output is not bound at start (index 1 is ())
%call-start = ((f32[1024]), (), s32[]) call-start(%input_buffer), to_apply=%foo

// Output is bound at update (index 1 becomes (f32[1024]))
%call-update = ((f32[1024]), (f32[1024]), s32[]) call-update(%call-start, %output_buffer)

%result = (f32[1024]) call-done(%call-update)

تجزیه‌کننده این را به صورت زیر تجزیه می‌کند:

%async-start = ((f32[1024]), (), s32[]) async-start(%input_buffer), calls=%foo
%async-update = ((f32[1024]), (f32[1024]), s32[]) async-update(%async-start, %output_buffer)
%result = (f32[1024]) async-done(%async-update)

مثال با async-done (بدون async-update )

اگر هیچ مرحله به‌روزرسانی میانی وجود نداشته باشد، خروجی می‌تواند مستقیماً به async-done متصل شود:

// Output is not bound at start (index 1 is ())
%call-start = ((f32[1024]), (), s32[]) call-start(%input_buffer), to_apply=%foo

// Output is bound at done
%result = (f32[1024]) call-done(%call-start)

تجزیه‌کننده این را به صورت زیر تجزیه می‌کند:

%async-start = ((f32[1024]), (), s32[]) async-start(%input_buffer), calls=%foo
%result = (f32[1024]) async-done(%async-start)

مثالی با async-update میانی و مقیدسازی خروجی در async-done

اگر مراحل به‌روزرسانی میانی برای اتصال عملوندها وجود داشته باشد، اما خروجی هنوز در انتها محدود شده باشد:

// Output is not bound at start, no operands bound
%call-start = ((), (), s32[]) call-start(), to_apply=%foo

// Operands are bound at update, but output remains unbound (index 1 is ())
%call-update = ((f32[], f32[]), (), s32[]) call-update(%call-start, %operand0, %operand1)

// Output is bound at done
%result = f32[] call-done(%call-update)

تجزیه‌کننده این را به صورت زیر تجزیه می‌کند:

%async-start = ((), (), s32[]) async-start(), calls=%foo
%async-update = ((f32[], f32[]), (), s32[]) async-update(%async-start, %operand0, %operand1)
%result = f32[] async-done(%async-update)