במסמך הזה מתואר ניתוח ההוספה לאינדקס של HLO, וניתן להשתמש בו באופן סימבולי מפות מחשוב של הוספה לאינדקס עבור פעולות HLO. מפת ההוספה לאינדקס היא פונקציה שממפה אינדקסים של ארגומנט אחד למדדים של זווית אחרת, למשל מדדים של HLO פלט של הוראה למדדים של קלט הוראות HLO או להפך.
דוגמה
לשידור מ-tensor<20xf32> אל tensor<10x20x30xf32>
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
מפת ההוספה לאינדקס מהפלט לקלט היא (i, j, k) -> (j) עבור i in
[0, 10], j in [0, 20] ו-k in [0, 30].
למה בחרנו לעשות זאת?
GPU XLA משתמש בכמה פתרונות מותאמים במיוחד כדי להתבסס על איחוד, אופרנד תזמון שימוש ואריחים (פרטים נוספים בהמשך). מטרת ההוספה לאינדקס מספק רכיב לשימוש חוזר בתרחישים לדוגמה כאלה. ניתוח הוספה לאינדקס מבוסס על תשתית Affine Map של MLIR ומוסיף סמנטיקה של HLO.
מיזוג
הסיבה למיזוג הזיכרון יכולה להיות אפשרית במקרים לא טריים, אנחנו יודעים אילו אלמנטים/פלחים של הקלט נקראים כדי לחשב רכיב הפלט.
ניצול אופרנדים
שימוש חיצוני ב-XLA מציין את מידת הקלט של כל קלט של ההוראה בהנחה שנעשה שימוש מלא בפלט שלו. בשלב זה, לא מתבצע חישוב של ניצול גם במקרה כללי. ניתוח ההוספה לאינדקס מאפשר לחשב את ניצול המשאבים בצורה מדויקת.
פריסה
ריבוע/פרוסה הם קבוצת משנה ריבועית-היפר של טינסור, שמוגדר באמצעות פרמטרים של סטיות, גדלים וצעדי סריקה. הפצת אריחים היא דרך לחשב פרמטרים של אריח היצרן/הצרכן של ההפעלה באמצעות הפרמטרים של פריסת הפריסה עצמה. יש כבר ספרייה דגימת softmax ונקודות. אפשר להפוך את ההפצה של המשבצות לגנרית וחזקה יותר אם היא תתבצע באמצעות מפות להוספה לאינדקס.
פונקציה ודומיין
מפת האינדקס היא פונקציה f(x) = f(d, r, rt)
שממפה אינדקס מרובה d של טינסור A לרכיבים/טווחים של
טינסור B. הפרמטר r מתייחס לטווחי האינדקסים של
המאפיינים שקיימים בטינזור B אבל לא בטינזור A.
הפרמטר rt מתייחס לערכים של סביבת זמן הריצה, למשל אינדקסים לפעולת איסוף
לדוגמה, אם יש ירידה מ-tensor<2x4x8x16xf32> ל-
tensor<4x8xf32>, אז מפת ההוספה לאינדקס מהפלט הדו-ממדי לקלט ב-4D היא
(d0, d1) -> (r0, d0, d1, r1), כאשר d_i הם משתני המאפיין
תואמים למדדים של ארגומנט הפלט. משתני הטווח r_j מקודדים
מספר ערכים, כלומר כדי לחשב רכיב (d0, d1) של הפלט, אנחנו צריכים
(r0, d0, d1, r1) רכיבי הקלט, כאשר r0 in [0, 1] ו-
r1 in [0, 15]
אפשר ליצור את המיפוי הזה לפי המאפיינים של הוראות HLO או אפשר ליצור מיפויים של הוראות שלא משולבות כדי לקבל הוספה לאינדקס עבור מיזוג. גם המיפוי כולל דומיין שמציין אילו רכיבים של הטנזור. המיפוי קיים.
f(x) s.t.
lb <= g(x) <= ub
מכיוון שאנחנו רוצים לצמצם את החישוב מחדש, אנחנו צריכים ספרייה לחישוב סימבולי. XLA כבר תלוי ב-MLIR, לכן אנחנו mlir::AffineMap במקום לכתוב ספריית אריתמטית סימבולית אחרת.
AffineMap טיפוסי נראה כך
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
ל-AffineMap יש שני סוגי פרמטרים: מאפיינים וסמלים.
המאפיינים תואמים למשתני המאפיין d, כאשר הסמלים תואמים ל-
משתני הטווח r ו-RT rt. AffineMap לא מכיל/ה אף אפשרות
מטא-נתונים לגבי טווחים של המאפיינים, לכן אנחנו צריכים לספק את הנתונים האלה
את עצמנו.
struct Interval {
int64_t lower;
int64_t upper;
};
// Dimension variable represents a dimension of a tensor or a GPU grid.
struct DimVar {
Interval bounds;
};
// RangeVar variable represents a range of values, e.g. to compute a single
// element of the reduction's result we need a range of values from the input
// tensor.
struct RangeVar {
Interval range;
};
// RTVar represents a runtime value, e.g. a dynamic offset in
// HLO dynamic-update-slice op.
struct RTVar {
Interval feasible_values;
const HloInstruction* hlo;
// This is a map from the iteration space of the corresponding indexing map to
// the iteration space of `hlo`. It shows what element of `hlo` we need to
// extract to get the runtime value for the RTVar.
mlir::AffineMap map;
};
class IndexingMap {
mlir::AffineMap affine_map_;
std::vector<DimVar> dim_vars_;
std::vector<RangeVar> range_vars_;
std::vector<RTVar> rt_vars_;
llvm::DenseMap<mlir::AffineExpr, Interval> constraints_;
};
dim_vars_ מקודדים את מגבלות התיבה כולל של המאפיין
d במפת ההוספה לאינדקס, שחופפים בדרך כלל
של טנזור הפלט עבור פעולות כמו שינוי, הקטנה, יסודות, נקודה,
יש כמה יוצאים מן הכלל,
HloConcatenateInstruction.
range_vars_ מקודד את הערכים האפשריים שפרמטרים מסוג r יכולים לקבל.
rt_vars_ מאחסן את הוראות ה-hlo המשויכות יחד עם דפוסי הגישה שלהן והערכים האפשריים בסביבת זמן הריצה. לדוגמה, ההיסט הוא דינמי
לHloDynamicSliceInstruction ביום אחד. ל-RTVar התואם יהיה HloInstruction* שיוצר טינסור מסדר 0 עם תבנית הגישה (d0) -> (), כי לכל רכיב של הפלט אנחנו מחלצים את אותו רכיב מטינסור ההיסט כדי לחשב את האינדקס של הקלט. אפשר גם להניח שההיסט של הפלחים תמיד בין 0 ל-tensor_size - slice_size - 1.
הדוגמה הבאה תעזור לכם להבין מה המשמעות של כל האפשרויות שלמעלה.
הוספת מפות לאינדקס עבור משימות ניהול לא מחוברות
אלמנט אחר אלמנט
בפעולות על רכיבים, מפת הוספת האינדקס היא זהות.
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
הפלט שמתקבל למפות קלט:
- פלט -> input_i:
(d0, d1) -> (d0, d1)
domain:
d0 in [0, 9]
d1 in [0, 19]
הקלט ליצירת מפות כפלט
- input_i -> output:
(d0, d1) -> (d0, d1)
domain:
d0 in [0, 9]
d1 in [0, 19]
שידור
המשמעות של שידור היא שחלק מהמאפיינים יוסרו כאשר נמפה כפלט לקלט, ונוסף כאשר אנחנו ממפים את הקלט לפלט.
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
הפלט למפת הקלט:
(d0, d1, d2) -> (d1)
domain:
d0 in [0, 9]
d1 in [0, 19]
d2 in [0, 29]
המיפוי של הקלט לפלט
(d0)[s0, s1] -> (s0, d0, s1)
domain:
d0 in [0, 19]
s0 in [0, 9]
s1 in [0, 29]
שימו לב שעכשיו יש לנו s בצד שמאל עבור הקלט לפלט
מיפוי. אלו הסמלים שמייצגים טווחי ערכים. לדוגמה, ב-
במקרה הספציפי הזה, כל רכיב של קלט עם האינדקס d0 ממופה
פרוסה בגודל 10x130 של הפלט.
קבוע ו-Iota
ל-URLs האלה אין פרמטרים של קלט, כך שאין מה לחשב כדי להוסיף אותם לאינדקס.
DynamicSlice
DynamicSlice הוא בדיוק כמו Slice, אבל ההיסטים הם דינמיים.
src = s32[2,2,258] parameter(0)
of1 = s32[] parameter(1)
of2 = s32[] parameter(2)
of3 = s32[] parameter(3)
ds = dynamic-slice(s32[2,2,258] src, s32[] of1, s32[] of2, s32[] of3), dynamic_slice_sizes={1, 2, 32}
הפלט שצריך להזין במפת הקלט של src:
(d0, d1, d2)[s0, s1, s2] -> (d0 + s0, d1 + s1, d2 + s2)
domain:
d0 in [0, 0]
d1 in [0, 1]
d2 in [0, 31]
s0 in [0, 1]
hlo: of1 = s32[] parameter(1)
(d0, d1, d2) -> ()
s1 in [0, 0]
hlo: of2 = s32[] parameter(2)
(d0, d1, d2) -> ()
s2 in [0, 226]
hlo: of3 = s32[] parameter(3)
(d0, d1, d2) -> ()
שימו לב שעכשיו יש לנו s בצד שמאל למיפוי הקלט לפלט.
אלה הסמלים שמייצגים ערכים של סביבת זמן ריצה. לדוגמה, במקרה הזה, לכל רכיב של הפלט עם האינדקסים d0, d1, d2 אנחנו ניגשים להיסטים של הפרוסות of1, of2 ו-of3 כדי לחשב את האינדקס של הקלט.
הפסקות הזמן של משתני זמן הריצה נובעות מהנחה שכל החלק נשאר בגבולות.
המפה של הפלט לקלט עבור of1, of2 ו-of3:
(d0, d1, d2) -> ()
domain:
d0 in [0, 0]
d1 in [0, 1]
d2 in [0, 31]
DynamicUpdateSlice
src = s32[20,30] parameter(0)
upd = s32[5,10] parameter(1)
of1 = s32[] parameter(2)
of2 = s32[] parameter(3)
dus = s32[20,30] dynamic-update-slice(
s32[20,30] src, s32[5,10] upd, s32[] of1, s32[] of2)
מפת הפלט לתשתית של src היא טריוויאלית. אפשר להגביל את הדומיין לאינדקסים שלא עודכנו כדי לקבל תוצאות מדויקות יותר, אבל בשלב הזה מפות ההוספה לאינדקס לא תומכות באילוצים של אי-שוויון.
(d0, d1) -> (d0, d1)
domain:
d0 in [0, 19]
d1 in [0, 29]
הפלט שצריך להזין במפת הקלט של upd:
(d0, d1)[s0, s1] -> (d0 - s0, d1 - s1)
domain:
d0 in [0, 19]
d1 in [0, 29]
s0 in [0, 15]
hlo: of1 = s32[] parameter(2)
(d0, d1) -> ()
s1 in [0, 20]
hlo: of2 = s32[] parameter(3)
(d0, d1) -> ()
שימו לב שנוספה עכשיו s בצד שמאל למיפוי הקלט-פלט.
אלה הסמלים שמייצגים ערכים בסביבת זמן הריצה. לדוגמה, במקרה הזה, לכל רכיב של הפלט עם האינדקסים d0, d1 אנחנו ניגשים להיסטים של הפרוסה of1 ו-of2 כדי לחשב את האינדקס של הקלט. האינטרוולים
של משתני זמן הריצה, נגזרים בהנחה שכל הפלח נשאר
גבולות.
הפלט שמתקבל למפת הקלט עבור of1 ו-of2:
(d0, d1) -> ()
domain:
d0 in [0, 19]
d1 in [0, 29]
איסוף
יש תמיכה רק ב-gather המפושט. אפשר לעיין ב-[gather_simplifier]. (https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/gather_simplifier.h).
operand = f32[33,76,70] parameter(0)
indices = s32[1806,2] parameter(1)
gather = f32[1806,7,8,4] gather(operand, indices),
offset_dims={1,2,3},
collapsed_slice_dims={},
start_index_map={0,1},
index_vector_dim=1,
slice_sizes={7,8,4}
המפה של הפלט לתשתית הקלט עבור operand:
(d0, d1, d2, d3)[s0, s1] -> (d1 + s0, d2 + s1, d3)
domain:
d0 in [0, 1805]
d1 in [0, 6]
d2 in [0, 7]
d3 in [0, 3]
s0 in [0, 26]
hlo: indices = s32[1806,2]{1,0} parameter(1)
(d0, d1, d2, d3) -> (d0, 0)
s1 in [0, 68]
hlo: indices = s32[1806,2]{1,0} parameter(1)
(d0, d1, d2, d3) -> (d0, 1)
שימו לב שעכשיו יש לנו s בצד שמאל למיפוי הקלט לפלט.
אלה הסמלים שמייצגים ערכים בסביבת זמן הריצה. לדוגמה, במקרה הזה, לכל רכיב של הפלט עם האינדקסים d0, d1, d2, d3 אנחנו מחלצים את הרכיבים (d0, 0) ו-(d0, 1) מהטנסור indices.
הפלט שצריך להזין במפת הקלט של indices:
(d0, d1, d2, d3)[s0] -> (d0, s0)
domain:
d0 in [0, 1805]
d1 in [0, 6]
d2 in [0, 7]
d3 in [0, 3]
s0 in [0, 1]
משתנה הטווח s0 מראה שאנחנו צריכים את כל השורה (d0, *) של
את Tensor indices כדי לחשב רכיב של הפלט.
החלפה
מפת ההוספה לאינדקס של העברה אוטומטית היא תמורה של מאפייני הקלט/פלט.
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
המיפוי של הפלט לקלט:
(d0, d1, d2, d3) -> (d0, d3, d1, d2)
domain:
d0 in [0, 2]
d1 in [0, 5]
d2 in [0, 127]
d3 in [0, 12287]
המיפוי של הקלט לפלט:
(d0, d1, d2, d3) -> (d0, d2, d3, d1)
domain:
d0 in [0, 2]
d1 in [0, 12287]
d2 in [0, 5]
d3 in [0, 127]
היפוך
מפת הוספה לאינדקס לשינויים הפוכים של המאפיינים שהוחזרו ל-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}
הפלט למפת הקלט:
(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3)
domain:
d0 in [0, 0]
d1 in [0, 16]
d2 in [0, 8]
d3 in [0, 8]
המיפוי של הקלט לפלט:
(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3)
domain:
d0 in [0, 0]
d1 in [0, 16]
d2 in [0, 8]
d3 in [0, 8]
(Variadic)Reduce
להפחתה משתנה יש מספר קלטים ומספר כניסות, המפה מפלט אל הקלט מוסיף את המימדים המצומצמים. אז היא מתנהגת כמו שידור במובן מסוים.
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=max
המיפוי של הפלט לקלט:
- פלט -> input_j:
(d0)[s0] -> (s0, d0)
domain:
d0 in [0, 9]
s0 in [0, 255]
- פלט -> init_j:
(d0) -> ()
domain:
d0 in [0, 9]
הקלט למיפוי הפלט:
- קלט_i -> output_j:
(d0, d1) -> (d1)
domain:
d0 in [0, 255]
d1 in [0, 9]
- init_i -> output_j:
()[s0] -> (s0)
domain:
s0 in [0, 9]
עבור i, j = 0, ... 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]}
המיפוי של הפלט לקלט:
(d0, d1, d2) -> (d0 + 5, d1 * 7 + 3, d2 * 2)
domain:
d0 in [0, 4]
d1 in [0, 2]
d2 in [0, 24]
המיפוי של הקלט לפלט:
(d0, d1, d2) -> (d0 - 5, (d1 - 3) floordiv 7, d2 floordiv 2)
domain:
d0 in [5, 9]
d1 in [3, 17]
d2 in [0, 48]
(d1 - 3) mod 7 in [0, 0]
d2 mod 2 in [0, 0]
שינוי הצורה
יש כמה סוגים של יצירת תבניות מחדש.
כיווץ הצורה
זוהי שינוי צורה 'לינאריזציה' ממאפיינים N-ממדיים למאפיין 1-ממדי.
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
הפלט למפת הקלט:
(d0) -> (d0 floordiv 8, d0 mod 8)
domain:
d0 in [0, 31]
מפת הקלט ליצירת פלט:
(d0, d1) -> (d0 * 8 + d1)
domain:
d0 in [0, 3]
d1 in [0, 7]
הרחבת הצורה
זוהי פעולה הפוכה של 'התכווצות צורה', שמשנה את הצורה של קלט חד-ממדי לפלט N-ממדי.
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
המיפוי של הפלט לקלט:
(d0, d1) -> (d0 * 8 + d1)
domain:
d0 in [0, 3]
d1 in [0, 7]
המיפוי של הקלט לפלט:
(d0) -> (d0 floordiv 8, d0 mod 8)
domain:
d0 in [0, 31]
שינוי צורה כללי
אלה פעולות שינוי הצורה שאי אפשר לייצג כצורה אחת של התרחבות או כוונון. אפשר לייצג אותם רק כקומפוזיציה של 2 או יותר צורות של התרחבות או כיווץ.
דוגמה 1: ליניאריזציה-ביטול ליניאריזציה.
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
אפשר לייצג את השינוי הזה כקומפוזיציה של צורת קריסת tensor<4x8xf32> ל-tensor<32xf32> ואז של התרחבות הצורה ל-tensor<2x4x4xf32>.
המיפוי של הפלט לקלט:
(d0, d1, d2) -> (d0 * 2 + d1 floordiv 2, d2 + (d1 mod 2) * 4)
domain:
d0 in [0, 1]
d1 in [0, 3]
d2 in [0, 3]
מפת הקלט ליצירת פלט:
(d0, d1) -> (d0 floordiv 2, d1 floordiv 4 + (d0 mod 2) * 2, d1 mod 4)
domain:
d0 in [0, 3]
d1 in [0, 7]
דוגמה 2: צורות משנה מורחבות ומכווצות
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
אפשר לייצג את הצורה הזו כהרכב של שני צורות מחדש. האפשרות הראשונה מכוונת את המאפיינים החיצוניים ביותר tensor<4x8x12xf32> אל tensor<32x12xf32>, והאפשרות השנייה מרחיבת את המאפיין הפנימי ביותר tensor<32x12xf32> אל tensor<32x3x4xf32>.
הפלט למפת הקלט:
(d0, d1, d2) -> (d0 floordiv 8, d0 mod 8, d1 * 4 + d2)
domain:
d0 in [0, 31]
d1 in [0, 2]
d2 in [0, 3]
מפת הקלט ליצירת פלט:
(d0, d1, d2) -> (d0 * 8 + d1, d2 floordiv 4, d2 mod 4)
domain:
d0 in [0, 3]
d1 in [0, 7]
d2 in [0, 11]
העברת נתונים (Bitcast)
אפשר לייצג פעולת bitcast בתור רצף של transpose-reshape-transpose. לכן, המפות של ההוספה לאינדקס הן רק הרכב של מפות הוספה לאינדקס עבור ברצף.
שרשור
מיפוי הפלט לקלט של concat מוגדר לכל הקלטים, אבל עם דומיינים לא חופפים, כלומר המערכת תשתמש רק באחד מהקלטים בכל פעם.
p0 = f32[2, 5, 7] parameter(0)
p1 = f32[2, 11, 7] parameter(1)
p2 = f32[2, 17, 7] parameter(2)
ROOT concat = f32[2, 33, 7] concatenate(f32[2, 5, 7] p0, f32[2, 11, 7] p1, f32[2, 17, 7] p2), dimensions={1}
הפלט אל הקלט ממפה:
- פלט -> קלט 1:
(d0, d1, d2) -> (d0, d1, d2)
domain:
d0 in [0, 1]
d1 in [0, 4]
d2 in [0, 6]
- פלט -> קלט 2:
(d0, d1, d2) -> (d0, d1 - 5, d2)
domain:
d0 in [0, 1]
d1 in [5, 15]
d2 in [0, 6]
- פלט -> קלט 3:
(d0, d1, d2) -> (d0, d1 - 16, d2)
domain:
d0 in [0, 1]
d1 in [16, 32]
d2 in [0, 6]
מקורות הקלט ליצירת מפות כפלט:
- קלט 1 -> פלט:
(d0, d1, d2) -> (d0, d1, d2)
domain:
d0 in [0, 1]
d1 in [0, 4]
d2 in [0, 6]
- קלט 2 -> פלט:
(d0, d1, d2) -> (d0, d1 + 5, d2)
domain:
d0 in [0, 1]
d1 in [0, 10]
d2 in [0, 6]
- קלט 3 -> פלט:
(d0, d1, d2) -> (d0, d1 + 16, d2)
domain:
d0 in [0, 1]
d1 in [0, 16]
d2 in [0, 6]
Dot
מפות ההוספה לאינדקס של נקודה דומות מאוד למפות ההפחתה.
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:
(d0, d1, d2)[s0] -> (d0, d1, s0)
domain:
d0 in [0, 3]
d1 in [0, 127]
d2 in [0, 63]
s0 in [0, 255]
- output -> input_2:
(d0, d1, d2)[s0] -> (d0, s0, d2)
domain:
d0 in [0, 3]
d1 in [0, 127]
d2 in [0, 63]
s0 in [0, 255]
מקורות הקלט ליצירת מפות כפלט:
- קלט_1 -> פלט:
(d0, d1, d2)[s0] -> (d0, d1, s0)
domain:
d0 in [0, 3]
d1 in [0, 127]
d2 in [0, 255]
s0 in [0, 63]
- input_2 -> output:
(d0, d1, d2)[s0] -> (d0, s0, d1)
domain:
d0 in [0, 3]
d1 in [0, 255]
d2 in [0, 63]
s0 in [0, 127]
פד
הוספה של PadOp לאינדקס היא ההפך מהוספה של SliceOp לאינדקס.
p0 = f32[4, 4] parameter(0)
p1 = f32[] parameter(1)
pad = f32[12, 16] pad(p0, p1), padding=1_4_1x4_8_0
הגדרת ה-padding 1_4_1x4_8_0 מציינת את lowPad_highPad_interiorPad_dim_0 x lowPad_highPad_interiorPad_dim_1.
הפלט שמתקבל למפות קלט:
- פלט -> קלט:
(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4)
domain:
d0 in [1, 7]
d1 in [4, 7]
(d0 - 1) mod 2 in [0, 0]
- פלט -> init:
(d0, d1) -> ()
domain:
d0 in [0, 11]
d1 in [0, 15]
ReduceWindow
התכונה DecreaseWindow ב-XLA מבצעת גם מרווח פנימי. לכן, אפשר לחשב את המפות להוספה לאינדקס כשיוצרים אותן משיוך של הוספה לאינדקס של ReduceWindow בלי הוספת שטח אחסון ושיוך של הוספה לאינדקס של PadOp.
c_inf = f32[] constant(-inf)
p0 = f32[1024, 514] parameter(0)
reduce-window = f32[1024, 3] reduce-window(p0, c_inf),
window={size=1x512 pad=0_0x0_0}, to_apply=max
הפלט שמתקבל למפות קלט:
- פלט -> קלט:
(d0, d1)[s0] -> (d0, d1 + s0)
domain:
d0 in [0, 1023]
d1 in [0, 2]
s0 in [0, 511]
- output -> init:
(d0, d1) -> ()
domain:
d0 in [0, 1023]
d1 in [0, 2]
הוספת מפות Google לאינדקס עבור Fusion
מפת אינדקס עבור Fusion op היא הרכב של מפות אינדקס עבור כל פעולה ב אשכול. יכול להיות שחלק מהקלט ייקרא כמה פעמים באמצעות דפוסי גישה.
קלט אחד, מספר מפות של הוספה לאינדקס
דוגמה ל-p0 + transpose(p0).
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 יהיו (d0, d1) -> (d0, d1) ו-(d0, d1) -> (d1, d0). כלומר, כדי לחשב רכיב אחד
מהפלט, יכול להיות שנצטרך לקרוא פעמיים את פרמטר הקלט.
קלט אחד, מפת ההוספה לאינדקס עם ביטול כפילויות
יש מקרים שבהם המפות להוספה לאינדקס הן למעשה זהות, למרות שזה לא ברור באופן מיידי.
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 היא פשוט (d0, d1, d2) -> (d2, d0, d1).
Softmax
יצירת האינדקס של פלט לקלט ממופה עבור parameter 0 עבור softmax:
(d0, d1, d2)[s0] -> (d0, d1, s0)
domain:
d0 in [0, 1]
d1 in [0, 64]
d2 in [0, 124]
s0 in [0, 124]
וגם
(d0, d1, d2) -> (d0, d1, d2)
domain:
d0 in [0, 1]
d1 in [0, 64]
d2 in [0, 124]
כאשר s0 מתייחס למאפיין הפנימי ביותר של הקלט.
הסבר פשוט על הוספת מפות לאינדקס
לכלי הפשטת ברירת המחדל של mlir::AffineMap ב-upstream אין אפשרות להניח הנחות לגבי הטווחים של המאפיינים או הסמלים. לכן, לא ניתן
מפשטים ביטויים באמצעות mod וdivביעילות.
אנחנו יכולים למנף את הידע שלנו לגבי הגבול התחתון והעליון של ביטויי משנה במפות Google חיבוריות כדי לפשט אותם עוד יותר.
הכלי הפשוט יכול לשכתב את הביטויים הבאים.
(d0, d1) -> (d0 + d1 floordiv 16, d1 mod 16)עבור d ב-[0, 6] x [0, 14]הופך ל-(d0, d1) -> (d0, d1)(d0, d1, d2) -> ((100d0 + 10d1 + d2) floorDiv 100, ((100d0 + 10d1 + d2) mod 100) floordiv 10, d2 mod 10)עבורdi in [0, 9]הופך ל-(d0, d1, d2) -> (d0, d1, d2).(d0, d1, d2) -> ((16d0 + 4d1 + d2) floordiv 8, (16d0 + 4d1 + d2) mod 8)עבורd_i in [0, 9]הופך ל-(d0, d1, d2) -> (2d0 + (4d1 + d2) floordiv 8,(4d1 + d2) mod 8).(d0, d1) -> (-(-11d0 - d1 + 109) floordiv 11 + 9)ל-י בעוד[0, 9] x [0, 10]הופך ל-(d0, d1) -> (d0).
כלי פשוט יותר להוספת מפות לאינדקס מאפשר לנו להבין שחלק יכולים ליצור מחדש ב-HLO זה את זה.
p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)
לאחר ההרכב של מפות לאינדקס והפישוט שלהן, נקבל
(d0, d1, d2) -> (d0, d1, d2).
הפשטת המפה להוספה לאינדקס מאפשרת גם לפשט את האילוצים.
- מגבלות מסוג
lower_bound <= affine_expr (floordiv, +, -, *) constant <= upper_boundהם נכתב מחדש כ-updated_lower_bound <= affine_expr <= updated_upped_bound. - אילוצים שתמיד מתקיימים, למשל
d0 + s0 in [0, 20]עבורd0 in [0, 5]ו-s0 in [1, 3], מוחקים. - ביטויים מדויקים במגבלות עוברים אופטימיזציה בתור ערך ההוספה לאינדקס מפה למעלה.
דוגמאות נוספות זמינות בכתובת indexing_map_test.cc.