צורות ופריסה

המבנה של פעולת XLA

דוגמה ל-HLO:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

המחרוזת הזו מורכבת מהרכיבים הבאים:

• שם הפעולה: add.936
  • זהו השם הייחודי של הפעולה.
• צורה: bf16[8,1,1280,16384]
  • זהו הפלט של הצורה של האופרטור. כאן, ה-dtype הוא bf16 והצורה היא [8,1,1280,16384].
• פריסה (עם סידור חלונות): 3,2,0,1:T(8,128)(2,1)
  • כאן מתואר איך המערך מאוחסן בזיכרון. ‫3,2,0,1 מציין את סדר הצירים בזיכרון (כלומר, סדר עמודות, סדר שורות וכו') ו-T(8,128)(2,1) מציין את הפסיפס והריפוד שנעשה בהם שימוש.
  • הפריסה היא אופציונלית. אם לא מציינים ערך, לא מתבצעת חלוקה למשבצות והמערכת מניחה שהמימדים מסודרים מהמימד הראשי ביותר למימד המשני ביותר.
• פעולה: add
  • הפעולה שמבוצעת. במקרה הזה, זהו Add, שמוזכר גם בשם הפעולה.
• ארגומנטים: exponential.183, broadcast.3115
  • הפעולה הזו מקבלת שני ארגומנטים, שצוינו עם השמות הייחודיים שלהם.

נבחן דוגמה נוספת, פעולת מיזוג:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

בנוסף לרכיבים שתוארו קודם, הוא כולל את הרכיבים הבאים:

• מאפיינים: kind ו-calls
  • הם מספקים מידע נוסף על הפעולה שמבוצעת, ובמקרה הזה: מיזוג.
• מיקום בזיכרון (מזהה של מקום בזיכרון): S(1)
  • הערך הזה מציין את המיקום בזיכרון שבו המערך מאוחסן. S(1) ‫here מציין שהמערך הזה נמצא ב-VMEM (ב-TPU).
• פרטים לגבי הצורה והפריסה של ארגומנט הקלט %fusion.32

בקטעים הבאים מוסבר על צורות, על פריסה ועל מזהים של מרחב זיכרון. מידע נוסף על סידור חלונות צמודים זמין במאמר בנושא פריסה של חלונות צמודים.

צורות

פרוטוקול XLA ShapeProto (xla_data.proto) מתאר את מספר המימדים, הגודל וסוג הנתונים של מערך N-ממדי (array בקיצור).

מונחים, סימון ומוסכמות

• המספר האמיתי של המאפיינים במערך הוא מספר המאפיינים שגודלם גדול מ-1.

• המאפיינים ממוספרים מ-0 עד N-1 עבור מערך N ממדי. הגודל של מאפיין הוא מספר שלם לא שלילי. בפרט, הגודל 0 הוא גודל תקין. מספרי המאפיינים הם תוויות שרירותיות לנוחותכם. הסדר של מספרי המאפיינים האלה לא מרמז על סדר משני/עיקרי מסוים בפריסה של הצורה. הפריסה נקבעת על ידי פרוטוקול LayoutProto.

• לפי המוסכמה, המאפיינים מופיעים בסדר עולה של מספר המאפיין. לדוגמה, במערך תלת-ממדי בגודל [A x B x C], הגודל של מימד 0 הוא A, הגודל של מימד 1 הוא B והגודל של מימד 2 הוא C.

  חלק מכלי השירות ב-XLA תומכים גם באינדקס שלילי כמו ב-Python: מימד ‎-1 הוא המימד האחרון (שווה ל-N-1 למערך N מימדי). לדוגמה, במערך התלת-ממדי שמתואר למעלה, המימד ‎-1 הוא בגודל C, המימד ‎-2 הוא בגודל B וכן הלאה.

• למערכים דו-ממדיים, תלת-ממדיים וארבע-ממדיים יש לרוב אותיות ספציפיות שמשויכות לממדים. לדוגמה, למערך דו-ממדי:

  • מאפיין 0: y
  • מאפיין 1: x

  למערך תלת-ממדי:

  • מאפיין 0: z
  • מאפיין 1: y
  • מאפיין 2: x

  עבור מערך 4D:

  • מאפיין 0: p
  • מאפיין 1: z
  • מאפיין 2: y
  • מאפיין 3: x

• פונקציות ב-XLA API שמקבלות מימדים עושות זאת בסדר עולה של מספר המימד. הסדר הזה תואם לסדר שבו משתמשים כשמעבירים מאפיינים כ-initializer_list, למשל:

  ShapeUtil::MakeShape(F32, {A, B, C, D})

  תיצור צורה שמערך גודל המימד שלה מורכב מהרצף [A, B, C, D].

פריסה

ב-proto של LayoutProto מתואר איך מערך מיוצג בזיכרון. הוא כולל את השדות הבאים:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

סדר המאפיינים ממאפיין משני למאפיין ראשי

שדה החובה היחיד הוא minor_to_major. בשדה הזה מתואר הסדר של המימדים בתוך צורה, מהקטן לגדול. הערכים ב-minor_to_major הם סדר של המימדים במערך (0 עד N-1 למערך N מימדי), כאשר הערך הראשון הוא המימד הכי קטן עד הערך האחרון שהוא המימד הכי גדול. המאפיין הכי משני הוא המאפיין שמשתנה הכי מהר כשעוברים בין הרכיבים של המערך שמוצג בזיכרון ליניארי.

לדוגמה, נניח שיש מערך דו-ממדי בגודל [2 x 3]:

a b c
d e f

בדוגמה הזו, מידה 0 היא מידה 2 ומידה 1 היא מידה 3. אם השדה minor_to_major בפריסה הוא [0, 1], אז מאפיין 0 הוא המאפיין המשני ביותר ומאפיין 1 הוא המאפיין הראשי ביותר. זה מתאים לפריסה הבאה בזיכרון ליניארי:

a d b e c f

סדר המימדים מהקטן לגדול, מ-0 עד N-1, דומה ל-column-major (למימדים דו-ממדיים). בהנחה של סדר מונוטוני של מאפיינים, דרך נוספת להתייחס לפריסה הזו בקוד היא פשוט 'מאפיין 0 הוא משני'.

לעומת זאת, אם השדה minor_to_major בפריסה הוא [1, 0], הפריסה בזיכרון ליניארי היא:

a b c d e f

סדר מימדים ממינורי למייג'ור של N-1 עד 0 למערך N מימדי דומה לסדר שורות (למערכים דו-מימדיים). בהנחה שיש סדר מונוטוני של מאפיינים, דרך נוספת שבה אנחנו עשויים להתייחס לפריסה הזו בקוד היא פשוט "מאפיין 0 הוא המאפיין הראשי".

סדר ברירת המחדל של מינור למייג'ור

פריסת ברירת המחדל של צורות שנוצרו לאחרונה היא 'סדר המאפיינים הוא מהגדול לקטן' (כלומר [N-1, ..., 0]).

מרווח

בשדה tail_padding_alignment_in_elements מגדירים את היישור של המערך tiled במונחים של מספר האלמנטים. אחרי שמחילים את האריחים, מוסיפים רכיבים עם ריווח בסוף הפריסה עד שמספר הרכיבים הכולל הוא כפולה של הערך הזה.

הוספה לאינדקס של מערכים

המחלקות IndexUtil ב-index_util.h מספקות כלי עזר להמרה בין אינדקסים רב-ממדיים לאינדקסים לינאריים, בהינתן צורה ופריסה. אינדקסים רב-ממדיים כוללים אינדקס int64 לכל מאפיין. אינדקסים ליניאריים הם ערך int64 יחיד שמציין את המיקום במאגר שמכיל את המערך. בספריות shape_util.h ו-layout_util.h אפשר למצוא כלי עזר שמפשטים את היצירה והשינוי של צורות ופריסות.

מזהים של מרחבי זיכרון

ב-HLO, כל מערך יכול להיות מתויג במזהה של מרחב זיכרון, שנכתב כ-S(n).

• ‫S(0) (לרוב מושמט) מציין זיכרון HBM (זיכרון ברוחב פס גבוה) של המכשיר.
• ‫S(1) מייצג זיכרון וירטואלי במכשיר (VMEM).
• ‫S(2), ‏S(3) וכו' תואמים למרחבי זיכרון נוספים שספציפיים למכשיר.
• ‫S(5) מציין את זיכרון המארח.