הוראות אסינכרוניות ל-HLO

1. הוספת פעולות אסינכרוניות ל-HLO היא מסורבלת (למשל all-reduce-start ו all-reduce-done).
2. ייתכן שפיצול ההתחלה והסיום לא מספיקים עבור חלק מהשימוש האסינכרוני במקרים שונים.

כדי להתמקד בחסרון הראשון, אנחנו מציעים להוסיף קבוצה אחת אחרונה של קודי פעולות אסינכרוניים: kAsyncStart, kAsyncUpdate ו-kAsyncDone. הרעיון היא ליצור קוד opcode אסינכרוני גנרי שיכול לעטוף כל הוראה של HLO. הפעולה שתבוצע בפועל תקודד באופן אסינכרוני באמצעות חישוב שנקרא שההוראה היא רק השורש שלו, בפרמטרים לקלט. טיפול במאגר הנתונים הזמני של הקלט/פלט בזמן ההפעלה ויצירת כינויים ניתן לשתף לאחר כל פעולה אסינכרונית. ההוראות להפעלה אסינכרונית הפלט יהיה משולש של אופרנדים של הקלט, ערכי פלט מצב הביניים שנחוץ בשביל 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 התואם כדי למצוא את החישוב שנקרא.

יש לשים לב גם שהרכיב הראשון ב-tuple של הפלט של async-start עם פרמטר אופרנד, כך שהמאגר הזמני נשאר פעיל עד לפחות להוראה אסינכרונית. באופן דומה, הכינויים של הרכיב השני עם הפלט של async-done, וגם השלישי הוא מצב ההקשר שמשמש למעקב פעולה אסינכרונית. הייצוג הזה תומך גם במספר מעבדי Tensor הקלט ו/או הפלט של הפעולה האסינכרונית והכינויים פועלים באופן זהה דרך:

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

תחביר סוכר

מכיוון שנעשה חישוב נפרד כדי להגדיר את הפעולה שתינתן באופן אסינכרוני הוא קצת מסורבל, אנחנו מציעים גם תחביר sugar להדפיס ולנתח פעולות אסינכרוניות באופן אוטומטי כאילו הן מהדרגה הראשונה מקודדים. הרעיון הוא לטפל בסיומות '-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)

כדי שלא ליצור אי בהירות, המאמת לא יאפשר פעולה: להיות מוקף ב-async-start אם הגדרנו במפורש opcode עם הסיומות ' -start' ו/או ' -done'. גם זו פעולת בריחה למקרה שיהיו לנו הוראות שמחייבות טיפול ברמת HLO, לא מתאים למודל המתואר למעלה (למשל, קלט/פלט עם כינויים מאגרי נתונים זמניים). לכן, בהתחלה, copy-start/copy-done, collective-permute-start/collective-permute-done וכו' ימשיכו להשתמש את פקודות השירות שלהם מהמחלקה הראשונה במקום async-start מתוך async-done פעולות פעולות עד שננקה את הקוד כדי להסיר אותן " -start"/"-done".