- افزودن عملیات async به HLO دشوار است (یعنی
all-reduce-start
وall-reduce-done
). - تقسیم شروع و انجام شده ممکن است برای برخی از موارد استفاده ناهمزمان ناکافی باشد.
برای هدف قرار دادن اولین نقص، پیشنهاد می کنیم آخرین مجموعه از کدهای عملیاتی ناهمزمان جدید را معرفی کنیم: 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” کد را بالا ببرید.