פריסת משבצות

איור 1

איור 1 מראה את הפריסה של המערך F32[3,5] בזיכרון באמצעות אריחים בגודל 2x2. צורה עם הפריסה הזו נכתבת כך: F32[3,5]{1,0:T(2,2)}, כאשר 1,0 מתייחס לסדר המימדים הפיזי (השדה minor_to_major בפריסה) ואילו (2,2) אחרי הנקודתיים מציינים משבצות של המידות הפיזיות באמצעות אריח 2x2.

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

האלמנטים הנוספים במרווח הפנימי לא חייבים להכיל ערך מסוים.

נוסחאות של אינדקס לינארי לריצוף בהתאם לצורה ולמשבצת

בלי אריחים, הרכיב e=(e_n, e_n-1, ... , e₁) במערך עם גבולות המערך d=(d_n, d_n-1, ... , d₁) (d1 הוא המאפיין המינורי ביותר) (d1 הוא המאפיין המשני ביותר) בפריסה לפי הסדר הראשי למשנהו במיקום:

linear_index(e, d)
= linear_index(e_n, e_n-1, ... , e₁), (d_n, d_n-1, ... , d₁))
= e_nd_n-1...d₁ + _n-1...d₁ + _n-1

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

כשמשתמשים בריצוף של גודל (t_n, t_n-1, ... , t₁), המערכת ממופה למיקום הזה בפריסה הסופית רכיב במערך עם אינדקסים (e_n, e_n-1, ... , e₁):

⌊, &index_with_tile(e, d, t)
= linear_index(⌊e/t⌋, e mod t), (⌈d/t⌉, t)) ( rithmetic is elementwise, (a,b) הוא שרשור

_n

אפשר לדמות את הפריסה כך שכוללת שני חלקים: (⌊e_n/t_n⌋, ... , ⌊e₁/t₁⌋), שתואם לאינדקס של קטעי המפה במערך של אריחי גודל (⌊e_n/t_n⌉, ... , ⌈d₁), ... , ⌈d₁. הפונקציה ceil מופיעה כאן: ⌈d_i/t_i⌉ כי אם המשבצות חורגות מגבולות המערך הגדול יותר, מוזן מרווח כמו באיור 1. האריחים וגם הרכיבים שבתוך האריחים מסודרים באופן רקורסיבי ללא אריחים.

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

– sitemap sitemap_with_tile (2,3), (3,5), (2,2))
= linear_index((1,1,0,1), (2,3,2,2))
= linear_index(1,1), (2,3)) 2
2 + linear_index(0,1), (2, 2 + linear_index(0,1), (2, 2)

אריחים בצבעי pad-reshape-transpos

פריסה המבוססת על אריחים פועלת באופן הבא:
שקול מערך של מאפיינים (d_n, d_n-1, ... , d1) (d1 הוא המאפיין המשני ביותר). אחרי הפריסה עם אריחים בגודל (t_n, t_n-1, ... , t₁) (t₁ הוא המאפיין המשני ביותר), אפשר לתאר את האריחים במונחים של שינוי צורה של רפידות באופן הבא.

1. המערך מתווסף ל-(⌈d_n/t_n⌉t_n, ... , ⌈d₁/t₁⌉t₁).
2. כל מימד i מחולק (⌈d_i/ti⌉, t_i), כלומר המערך מעוצב מחדש ל-
   (⌈d_n/t_n⌉, t_n, ... , ⌈d₁/t₁⌉, t₁).
   אין שינוי פיזי בפריסה הפיזית, ולכן השינוי הזה הוא רק חלק מפורמט ה-bitcast. אם לא כותבים באופן מפורש על ריצוף, אפשר לבטא מחדש כל צורה עם אותו מספר רכיבים כמו הצורה המרופדת. בדוגמה הזו רואים איך לבטא אריח בדרך הזו.
3. שינוי מיקום מתבצע על ידי הזזת t_n, ... , t₁ לממדים המשניים ביותר תוך שמירה על הסדר היחסי שלהם, כך שסדר המאפיינים מהמז'ורי ביותר לרוב המינורי הופך ל-
   (⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉, ... n₁⌉, ... n_).

הצורה הסופית כוללת את הקידומת
(⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉), שמתארת את מספר המשבצות בכל ממד. רכיב במערך (e_n, ... , e₁) ממופה לרכיב הזה בצורה הסופית:
(⌊e_n/t_n⌋, ... , ⌊e₀/t₀⌋, e_n mod t_n, ... , e₁). ניתן לראות בקלות שהאינדקס הלינארי של הרכיב תואם לנוסחה שלמעלה, כמצופה.

אריחים חוזרים

הריצוף של XLA הופך לגמיש עוד יותר אם מחילים אותו שוב ושוב.

איור 2

איור 2 מציג את האריחים בגודל 4x8 באמצעות שתי רמות אריחים (קודם 2x4 ואז 2x1). אנחנו מייצגים את החלוקה לאריחים חוזרים באופן הבא: (2,4)(2,1). כל צבע מציין אריח בגודל 2x4 וכל תיבת גבולות אדומה היא אריח בגודל 2x1. המספרים מציינים את האינדקס הלינארי בזיכרון של הרכיב הזה בפורמט של משבצות. הפורמט הזה תואם לפורמט של BF16 ב-TPU, אלא שהאריח הראשוני גדול יותר. כלומר, הריצוף הוא (8,128)(2,1), שבו המטרה של הריצוף השני פי 2x1 היא לאסוף יחד שני ערכי 16 ביט כדי ליצור ערך אחד של 32 ביט באופן שיתאים לארכיטקטורה של TPU.

שימו לב משבצת שנייה או מאוחרת יותר יכולה להתייחס לשני המימדים המשניים בתוך האריחים, שפשוט מסדר מחדש את הנתונים בתוך האריח, כמו בדוגמה הזו (8,128)(2,1), אבל יכול להתייחס גם למאפיינים העיקריים של כל האריחים באריחים הקודמים.

שילוב מימדים באמצעות אריחים

גם ריצוף XLA תומך בשילוב מידות. לדוגמה, הוא יכול לשלב את המאפיינים ב-F32[2,7,8,11,10]{4,3,2,1,0} ב-F32[112,110]{1,0} לפני הריצוף שלו באמצעות (2,3). האריח שבו נעשה שימוש הוא (עקבו אחרי, 1, או 1, ←3). בדוגמה הזו, כוכבית במשבצת מרמזת על כך שהיא לוקחת את המימד הזה ומשלבת אותה עם המאפיין השולי הבא. אפשר לכלול כמה מאפיינים סמוכים במאפיין אחד. מימד משויך מיוצג על ידי ערך אריח של -1 במימד הזה של האריח, שאינו תקף בדרך אחרת במשבצת כגודל מימד.

ליתר דיוק, אם מסירים את מימד i של הצורה באמצעות כוכבית באריח, לפני החלת ההגדרה הקודמת של אריחים, המימד מוסר גם מהצורה שעליה נחתם אריחים וגם מהווקטור האריח, והמערך של מימד i-1 מהצורה קושר מ-d_i-1 ל-_id_i-1. השלב הזה חוזר על עצמו עבור כל כוכבית בווקטור המשבצת.