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

  1. افزودن عملیات async به 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 نام مستعار دارد و عنصر سوم حالت متنی است که برای پیگیری عملیات ناهمزمان استفاده می‌شود. این نمایش همچنین از تانسورهای متعدد در ورودی و/یا خروجی عملیات ناهمزمان پشتیبانی می‌کند و aliasing به همین ترتیب عمل می‌کند:

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

قند نحوی

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

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

برای اینکه ابهام ایجاد نشود، اگر ما صریحاً کد عملیاتی را با پسوندهای "-start" و/یا "-done" برای آن عملیات تعریف کرده باشیم، تأیید کننده اجازه نخواهد داد که عملیات با async-start پیچیده شود. این همچنین یک دریچه فرار است در صورتی که دستورالعملی داشته باشیم که نیاز به درمان در سطح HLO داشته باشد که در مدل توضیح داده شده در بالا مناسب نباشد (مثلاً بافرهای ورودی/خروجی aliasing). بنابراین، در ابتدا، copy-start / copy-done , collective-permute-start / collective-permute-done و غیره به استفاده از کدهای opcode درجه یک مربوطه خود به جای کدهای opcode async-start / async-done جدید ادامه می دهند تا زمانی که ما پاک کنیم. برای حذف این کدهای عملیاتی “-start”/”-done” کد را بالا ببرید.