במסמך הזה נתאר את ניתוח האינדקס של HLO, שבו תוכלו לחשב באופן סימבולי מפות אינדקס עבור פעולות HLO. מפת האינדקס היא פונקציה שממפה אינדקסים של Tensor אחד לאינדקסים של אחר, למשל אינדקסים של פלט הוראה של HLO לאינדקסים של הוראות HLO או להיפך.
דוגמה
לשידור מ-tensor<20xf32> אל tensor<10x20x30xf32>
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
מפת האינדקס מהפלט לקלט היא \((i, j, k) \mapsto (j)\) עבור $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$.
למה בחרנו לעשות זאת?
XLA GPU משתמש בכמה פתרונות מותאמים אישית לתיאום, לשימוש באופרנד ולאריחים (פרטים נוספים בהמשך). המטרה של ניתוח ההוספה לאינדקס היא לספק רכיב לשימוש חוזר בתרחישים כאלה. ניתוח ההוספה לאינדקס מבוסס על התשתית של Affine Map ב-MLIR, ומוסיף את הסמנטיקה של HLO.
הצטברות
הסיבה לגבי צימוד הזיכרון יכולה להיות אפשרית במקרים לא טריוויאליים, כשאנחנו יודעים אילו רכיבים או פרוסות של הקלט נקראים כדי לחשב רכיב של הפלט.
שימוש של אופרנד
השימוש של עמודי ההוראה ב-XLSA מציין את מידת השימוש בכל קלט של ההוראה, בהנחה שנעשה בו שימוש מלא בפלט. נכון לעכשיו, השימוש לא מחושב במקרה של בקשות כלליות. ניתוח ההוספה לאינדקס מאפשר לחשב את השימוש בצורה מדויקת.
פריסה
אריח/פרוסה היא קבוצת משנה של טנסטור היפר-מלבני, עם פרמטרים של היסט, גדלים וצעדים. הפצת משבצות היא דרך לחשב פרמטרים של משבצות של המפעיל/הצרכן של ההפעלה באמצעות הפרמטרים של המשבצות של ההפעלה עצמה. כבר יש ספרייה שעושה את זה בשביל softmax ו-dot. הפצת האריחים יכולה להיות כללית ואמינה יותר אם היא מתבטאת באמצעות מפות של הוספה לאינדקס.
פונקציה ודומיין
מפת האינדקס היא פונקציה \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\) למיפוי אינדקסים מרובים \(\boldsymbol{d}\) של tensor \(A\) לרכיבים/טווחים של tensor \(B\). הפרמטר \(\boldsymbol{s}\) מתייחס לטווחי האינדקסים של המאפיינים שנמצאים ב-tensor \(B\), אבל לא ב-tensor \(A\).
לדוגמה, במקרה שיש לנו הפחתה מ-tensor<2x4x8x16xf32>
ל-tensor<4x8xf32>, מפת האינדקס מהפלט הדו-ממדי לקלט ה-4D היא
\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\), כאשר \(d_i\) הם הפרמטרים
של המאפיינים שתואמים לאינדקסים של Tensor הפלט. פרמטרים \(s_j\)מקודדים ערכים מרובים. כלומר, כדי לחשב \((d_0, d_1)\) רכיב של הפלט, אנחנו צריכים \((s_0, d_0, d_1, s_1)\) אלמנטים של הקלט, כאשר \(s_0 \in [0, 2)\) וגם
\(s_1 \in [0, 16)\).
אפשר לבנות את המיפוי הזה לפי המאפיינים של הוראות HLO או להרכיב את המיפויים של הוראות לא משולבות כדי ליצור אינדקס למיזוג. למיפוי יש גם דומיין, שמציין לאילו רכיבים ב-tensor יש מיפוי.
\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]
מכיוון שאנחנו רוצים למזער את החישוב מחדש, אנחנו צריכים ספרייה לחישובים סימבוליים. הקוד XLA תלוי כבר ב-MLIR, ולכן אנחנו משתמשים ב-mlir::AffineMap במקום לכתוב ספרייה סימבולית אריתמטית.
AffineMap טיפוסי נראה כך
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
ל-AffineMap יש בקלות שני סוגי פרמטרים: מאפיינים וסמלים שבהם נוכל להשתמש \(\boldsymbol d\) וגם \(\boldsymbol s\) בהתאמה.
המאפיין AffineMap לא מכיל מטא-נתונים לגבי טווחי המאפיינים, ולכן אנחנו צריכים לספק את הנתונים האלה בעצמנו.
struct Range {
int64_t lower_bound;
int64_t upper_bound;
};
struct IndexingMap {
mlir::AffineMap affine_map;
std::vector<Range> dimension_ranges;
std::vector<Range> symbol_ranges;
llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};
dim_ranges מקודדת את אילוצי התיבה הכוללים עבור הפרמטרים של המאפיינים \(\boldsymbol{d}\) במפת האינדקס, שחופפים בדרך כלל לצורה של טנזור הפלט עבור פעולות כמו טרנספוזיציה, הפחתה, רכיב, נקודה, אבל יש כמה יוצאים מן הכלל כמו HloConcatenateInstruction.
symbol_ranges מקודדים ערכים אפשריים \(\boldsymbol {s}\) שפרמטרים יכולים לקבל.
נבחן כל דוגמה ונבין את המשמעות של כל זה.
יצירת אינדקס של מפות לתפעול לא משולב
ברמת הרכיבים
עבור פעולות הרכיבים, מפת האינדקס היא זהות.
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
הפלט שיש להזין ממופה:
- פלט ->input_0: \((d_0, d_1) \mapsto (d_0, d_1)\) עבור $\boldsymbol{d} \in [0,9] \times [0, 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}(output)$
- פלט ->input_1: \((d_0, d_1) \mapsto (d_0, d_1)\) עבור $\boldsymbol{d} \in {\rm Dom} (output)$
הקלט למפות פלט
- קלט_i -> פלט: \((d_0, d_1) \mapsto (d_0, d_1)\) עבור $\boldsymbol{d} \in {\rm Dom}(output)$
להודיע לכולם
המשמעות של שידור היא שחלק מהמאפיינים יוסרו אם נמפה פלט לקלט ונוסיף אותו כשנמפה קלט לפלט.
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
הפלט למיפוי הקלט:
- פלט -> קלט: \((d_0, d_1, d_2) \mapsto (d_1)\) עבור $\boldsymbol{d} \in {\rm Dom}(output)$
הקלט למפת פלט
- קלט -> פלט: \((d_0) \mapsto (s_0, d_1, s_1)\) עבור $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9] \times [0, 29]$.
שים לב שעכשיו יש לנו את \(\boldsymbol s\) בצד שמאל למיפוי קלט לפלט. אלה הם הסמלים שמייצגים טווחי ערכים. לדוגמה, במקרה הספציפי הזה, כל רכיב של קלט עם אינדקס \(d_0\) ממופה לפרוסה 10x1x30 של הפלט.
קונסטנט ויוטה
בדרך כלל אין בהם פרמטרים של קלט, כך שאין מה לחשב עבורם אינדקס.
החלפה
מפת ההוספה לאינדקס של טרנספוזיציה היא תמורה של מאפייני קלט/פלט.
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
הפלט למיפוי הקלט:
- פלט -> קלט: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\) עבור \(\boldsymbol{d} \in {\rm Dom}(output)\)
הקלט למפת הפלט:
- קלט -> פלט: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\) עבור \(\boldsymbol{d} \in {\rm Dom}(input)\)
הפוך
מפת ההוספה לאינדקס לביצוע שינויים הפוכים בממדים ששוחזרו ל-$upper_bound(d_i) - d_i$:
p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}
הפלט למיפוי הקלט:
- פלט -> קלט: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(output)$
הקלט למפת הפלט:
- קלט -> פלט: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(input)$
(Variadic)צמצום
להפחתה הוראיאדית יש כמה מקורות קלט וכמה כניסות, המפה מהפלט לקלט מוסיפה את הממדים המצומצמים. כך, הוא מתנהג כמו היפוך לשידור במובן מסוים.
p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
dimensions={0}, to_apply=min
הפלט שיש להזין ממופה:
- פלט ->input_j: \((d_0) \mapsto (s_0, d_0)\) עבור $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9]$
- פלט -> init_j: \((d_0) \mapsto ()\) עבור $\boldsymbol{d} \in {\rm Dom}(output)$
הקלט לפלט ממופה:
- קלט_i ->output_j: \((d_0, d_1) \mapsto (d_1)\) עבור $\boldsymbol{d} \in {\rm Dom}(input)$
- init_i ->output_j: \(() \mapsto (s_0)\) עבור \(\boldsymbol{s} \in [0, 9]\)
של \(i, j = 0, \ldots, INPUT\\_COUNT\).
תבנית של תוכן דינמי
הוספה לאינדקס מהפלט לקלט של הפרוסה תוביל ליצירת מפת אינדקס מקוצרת, שחוקית לכל רכיב של הפלט. המיפוי מהקלט לפלט מוגבל לטווח מצומצם של הרכיבים של הקלט.
p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
slice={[5:10:1], [3:20:7], [0:50:2]}
הפלט למיפוי הקלט:
- פלט -> קלט: \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\) עבור \(\boldsymbol{d} \in {\rm Dom}(output)\)
הקלט למפת הפלט:
- קלט -> פלט: \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) עבור \(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) עם צעדים $[1, 7, 2]$.
TBD: יצירת אינדקס מקלט לפלט
שינוי הצורה
יש צורות שונות בטעמים שונים.
כיווץ הצורה
זוהי צורה מחדש "ליניארית" מ-N-D ל-1D.
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
הפלט למיפוי הקלט:
- פלט -> קלט: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) עבור $\boldsymbol{d} \in {\rm Dom}(output)$
הקלט למפת הפלט:
- קלט -> פלט: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) עבור $\boldsymbol{d} \in {\rm Dom}(input)$.
הרחבת הצורה
זוהי פעולה הפוכה של 'כיווץ צורה', והיא משנה קלט דו-ממדי לפלט N-D.
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
הפלט למיפוי הקלט:
- פלט -> קלט: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) עבור $\boldsymbol{d} \in {\rm Dom}(output)$
הקלט למפת הפלט:
- קלט -> פלט: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) עבור $\boldsymbol{d} \in {\rm Dom}(input)$.
שינוי צורה כללי
אלה תפעולי העיצוב מחדש שלא ניתן לייצג כצורת הרחבה או כיווץ יחידה. אפשר לייצג אותן רק כקומפוזיציה של 2 או יותר צורות הרחבה או כיווץ.
דוגמה 1: לינאריזציה-ביטול לינארי
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
ניתן לייצג את הצורה מחדש כקומפוזיציה של צורת כיווץ של tensor<4x8xf32> עד tensor<32xf32> ולאחר מכן כהרחבת צורה ל-tensor<2x4x4xf32>.
הפלט למיפוי הקלט:
- פלט -> קלט: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$
ל- \(\boldsymbol{d} \in {\rm Dom}(output)\)
הקלט למפת הפלט:
- קלט -> פלט: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$
של \(\boldsymbol{d} \in {\rm Dom}(input)\).
דוגמה 2: צורות משנה מורחבות ומכווצות
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
אפשר לייצג את הצורה מחדש כהרכב של שני צורות מחדש. המאפיין הראשון מכווץ את המאפיינים החיצוניים tensor<4x8x12xf32> ל-tensor<32x12xf32>, והמאפיין השני מרחיב את המאפיין הפנימי tensor<32x12xf32> לתוך tensor<32x3x4xf32>.
הפלט למיפוי הקלט:
- פלט -> קלט: \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\) עבור \(\boldsymbol{d} \in {\rm Dom}(output)\)
הקלט למפת הפלט:
- קלט -> פלט: \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\) עבור \(\boldsymbol{d} \in {\rm Dom}(input)\).
Bitcast
אפשר לייצג פעולת ביטcast כרצף של Transpose-reshape-transpose. לכן, מפות האינדקס שלו הן רק הרכב של מפות האינדקס לרצף הזה.
שרשור
מיפוי פלט לקלט של concat מוגדר לכל ערכי הקלט, אבל כשמדובר בדומיינים שאינם חופפים, כלומר ייעשה שימוש רק באחד ממקורות הקלט בכל פעם.
p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
dimensions={1}
הפלט למיפוי הקלט:
- פלט -> קלט 1:
\((d_0, d_1) \mapsto (d_0, d_1)\) עבור \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)
- פלט -> קלט 2:
\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) עבור $\boldsymbol{d} \in [0, 2] \times [50, 79]$
ערכי הקלט למפת הפלט:
- קלט 1 -> פלט: \((d_0, d_1) \mapsto (d_0, d_1)\) עבור $\boldsymbol{d} \in {\rm Dom}(input_1)$.
- קלט 2 -> פלט: \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) עבור $\boldsymbol{d} \in {\rm Dom}(input_2)$.
נקודה (פלט לקלט מיושם
מפות אינדקס עבור נקודה דומות מאוד למפות של צמצום.
p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
lhs_batch_dims={0}, rhs_batch_dims={0},
lhs_contracting_dims={2}, rhs_contracting_dims={1}
הפלט שמתקבל מהקלט ממופה באופן הבא:
- פלט ->input_1: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) עבור \(\boldsymbol{d} \in {\rm Dom}(output)\) ו- \(\boldsymbol{s} \in [0, 255]\)
- פלט ->input_2: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\) עבור \(\boldsymbol{d} \in {\rm Dom}(output)\) ו- \(\boldsymbol{s} \in [0, 255]\)
הקלט למפות הפלט:
- קלט_1 -> פלט: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) עבור \(\boldsymbol{d} \in {\rm Dom}(input_1)\) ו- \(\boldsymbol{s} \in [0, 63]\)
- קלט_2 -> פלט: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) עבור \(\boldsymbol{d} \in {\rm Dom}(input_2)\) ו- \(\boldsymbol{s} \in [0, 127]\)
הקטנת חלונות (TBD)
רפידה (TBD)
יצירת אינדקס של מפות עבור Fusion
מפת אינדקס של Fusion op היא יצירה של יצירת אינדקס של מפות לכל פעולה באשכול. יכול לקרות מצב שבו מקורות קלט מסוימים נקראים כמה פעמים עם דפוסי גישה שונים.
קלט אחד, כמה מפות של אינדקס
הנה דוגמה של \(p_0 + p_0^T\)
f {
p0 = f32[1000, 1000] parameter(0)
transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}
מפות האינדקס של פלט לקלט עבור p0 יהיו $(d_0, d_1) \mapsto (d_0,
d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$. המשמעות היא שכדי לחשב רכיב אחד בפלט ייתכן שנצטרך לקרוא פעמיים את פרמטר הקלט.
קלט אחד, מפת אינדקס ללא כפילויות
יש מקרים שבהם מפות האינדקס למעשה זהות, למרות שזה לא ברור מיד.
f {
p0 = f32[20, 10, 50] parameter(0)
lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}
מפת האינדקס של פלט לקלט עבור p0 במקרה הזה היא רק $(d_0, d_1, d_2)
\mapsto (d_2, d_0, d_1)$.
softmax
הוספת האינדקס מפלט לקלט ממופה ל-parameter 0 עבור softmax:
- \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
- \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)
\(\boldsymbol{d} \in {\rm Dom}(output)\) ו \(\boldsymbol{s} \in [0, 124]\) מתייחס למאפיין הפנימי ביותר של הקלט.
מפשט מפת האינדקס
מפשט ברירת המחדל של mlir::AffineMap ב-upstream לא יכול להניח הנחות לגבי טווחי המימדים או הסמלים. לכן היא לא יכולה לפשט ביטויים עם mod ו-div ביעילות.
אפשר למנף את הידע לגבי הגבול התחתון והעליון של ביטויי המשנה במפות האפיניות כדי לפשט אותם עוד יותר.
המפשט יכול לשכתב את הביטויים הבאים.
- \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) עבור $\boldsymbol{d} \in [0, 6] \times [0, 14]\( becomes \)(d_0, d_1) \mapsto (d_0, d_1)$
- $(d_0, d_1, d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_2 \mod 10)\( for \)__i, \mod 10)\( for \)d_i)\( becomes \)
- $(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod 8)\( for \)d_i \in [0, 9]\( becomes \)(d_0, d_1, +_4) +1_0, +4_2) +1_0.
- \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) עבור $\boldsymbol{d} \in [0, 9] \times [0, 10]\( becomes \)(d_0, d_1) \mapsto (d_0)$.
הכלי שמפשט מפת ההוספה לאינדקס מאפשר לנו להבין שחלק מהצורות המשורשרות ב-HLO מבטלים זה את זה.
p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)
לאחר ההרכבה של מפות האינדקס והפשטות שלהן, נקבל
\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).