ניתוח הוספה לאינדקס

ניתוח האינדוקס של 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].

למה בחרנו לעשות זאת?

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

מיזוג

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

ניצול אופרנדים

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

פריסה

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

מפת אינדקס

מפת אינדוקס היא שילוב של

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

הארגומנטים של הפונקציות מחולקים ל-3 קטגוריות כדי להסביר טוב יותר את המהות שלהם:

משתני המאפיין של הטנזור A או של רשת ה-GPU שאנחנו ממפים ממנה; הערכים ידועים באופן סטטי. רכיבי אינדקס נקראים גם משתני מאפיין.
משתני range. הם מגדירים מיפוי של אחד לרבים ומציינים קבוצה של רכיבים ב-B שמשמשים לחישוב ערך יחיד של A. הערכים ידועים באופן סטטי. המאפיין המצטמצם של כפל מטריצות הוא דוגמה למשתנה טווח.
משתני זמן ריצה שמוכרים רק במהלך ההרצה. לדוגמה, ארגומנט indices של פעולת gather.

התוצאה של הפונקציה היא אינדקס של טנסור היעד B.

בקיצור, פונקציית אינדוקס מטנזור A לטנזור B לפעולה x היא

map_ab(index in A, range variables, runtime variables) -> index in B.

כדי להפריד טוב יותר בין סוגי ארגומנטי המיפוי, אנחנו כותבים אותם כך:

map_ab(index in A)[range variables]{runtime variables} -> (index in B)

לדוגמה, נבחן את מפות האינדוקס של פעולת ה-reduce f32[4, 8] out = reduce(f32[2, 4, 8, 16] in, 0), dimensions={0,3}:

כדי למפות אלמנטים של in ל-out, אפשר לבטא את הפונקציה שלנו כך: (d0, d1, d2, d3) -> (d1, d2). המגבלות של המשתנים d0 in [0, 1], d1 in [0, 3], d2 in [0, 7], d3 in [0, 15] מוגדרות על ידי הצורה של in.
כדי למפות אלמנטים של out ל-in: ל-out יש רק שני מאפיינים, והפחתה יוצרת שני משתני טווח שמכסים את המאפיינים המופחתים. לכן פונקציית המיפוי היא (d0, d1)[s0, s1] -> (s0, d0, d1, s1), כאשר (d0, d1) הוא האינדקס של out. ‫s0, s1 הם טווחים שמוגדרים על ידי הסמנטיקה של הפעולה ומשתרעים על מימד 0 ו-3 של טנזור in. המגבלות הן d0 in [0, 3], d1 in [0, 7], s0 in [0,1], s1 in [0, 15].

חשוב לציין שברוב התרחישים אנחנו מתעניינים במיפוי מהרכיבים של הפלט. לחישוב

C = op1(A, B)
E = op2(C, D)

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

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

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

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

הטמעה

כדי לצמצם את הצורך בחישובים מחדש, אנחנו צריכים ספרייה לחישובים סימבוליים. XLA כבר תלויה ב-MLIR, ולכן אנחנו משתמשים ב-mlir::AffineMap במקום לכתוב עוד ספרייה של חשבון סמלי.

דוגמה ל-AffineMap

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

ב-AffineMap יש שני סוגים של פרמטרים: מאפיינים וסמלים. המאפיינים תואמים למשתני המאפיינים d, הסמלים תואמים למשתני הטווח r ולמשתני זמן הריצה rt. הקובץ AffineMap לא מכיל מטא-נתונים לגבי אילוצים של הפרמטרים, ולכן אנחנו צריכים לספק אותם בנפרד.

struct Interval {
 int64_t lower;
 int64_t upper;
};

class IndexingMap {
   // Variable represents dimension, range or runtime variable.
  struct Variable {
    Interval bounds;
    // Name of the variable is used for nicer printing.
    std::string name = "";
  };

  mlir::AffineMap affine_map_;

  // DimVars represent dimensions of a tensor or of a GPU grid.
  std::vector<Variable> dim_vars_;

  // RangeVars represent ranges of values, e.g. to compute a single element of
  // the reduction's result we need a range of values from the input tensor.
  std::vector<Variable> range_vars_;

  // RTVars represent runtime values, e.g. a dynamic offset in
  // HLO dynamic-update-slice op.
  std::vector<Variable> rt_vars_;
  llvm::DenseMap<mlir::AffineExpr, Interval> constraints_;
};

‫dim_vars_ קידוד האילוצים של התיבה כולל עבור משתני המימד d של מפת האינדקס, שבדרך כלל חופפים לצורה של טנסור הפלט עבור פעולות כמו transpose, ‏ reduce, ‏ elementwise, ‏ dot, אבל יש כמה חריגים כמו HloConcatenateInstruction.

range_vars_ כל הערכים שמשתני הטווח s מקבלים. משתני הטווח נדרשים כשצריך כמה ערכים כדי לחשב רכיב יחיד של הטנזור שממנו אנחנו ממפים, למשל, למיפוי אינדקסים של פלט->קלט של צמצומים או למיפוי קלט->פלט של שידורים.

rt_vars_ קידוד הערכים האפשריים בזמן הריצה. לדוגמה, ההיסט דינמי עבור 1D HloDynamicSliceInstruction. הערכים האפשריים של RTVar הם בין 0 ל-tensor_size - slice_size - 1.

constraints_ לכידת קשרים בין ערכים בטופס <expression> in <range>, למשל d0 + s0 in [0, 20]. יחד עם Variable.bounds הם מגדירים את ה'תחום' של פונקציית האינדקס.

כדי להבין את המשמעות של כל מה שצוין למעלה, נשתמש בדוגמה.

אינדוקס של מפות לפעולות לא משולבות

Elementwise

עבור פעולות ברמת הרכיב, מיפוי האינדקסים הוא זהות.

  p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  output = f32[10, 20] add(p0, p1)

מיפוי הפלט לקלט output -> p0:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

מיפוי הקלט לפלט p0 -> 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 ממופה לפרוסת פלט בגודל 10x1x30.

Iota

ל-Iota אין אופרנד טנסור קלט, ולכן אין ארגומנטים של אינדקס קלט.

iota = f32[2,4] iota(), dimensions={1}

פלט למיפוי קלט:

(d0, d1) -> ()
domain:
d0 in [0, 1]
d1 in [0, 3]

מיפוי קלט לפלט:

()[s0, s1] -> (s0, s1)
domain:
s0 in [0, 1]
s1 in [0, 3]

DynamicSlice

ל-DynamicSlice יש היסטים שמוכרים רק בזמן הריצה.

src = s32[2, 2, 258] parameter(0)
of1 = s32[] parameter(1)
of2 = s32[] parameter(2)
of3 = s32[] parameter(3)
ds = s32[1, 2, 32] dynamic-slice(src, of1, of2, of3), dynamic_slice_sizes={1, 2, 32}

מיפוי הפלט לקלט מ-ds אל src:

(d0, d1, d2){rt0, rt1, rt2} -> (d0 + rt0, d1 + rt1, d2 + rt2),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31],
rt0 in [0, 1],
rt1 in [0, 0],
rt2 in [0, 226]

שימו לב שעכשיו יש לנו rt בצד שמאל למיפוי של הקלט לפלט. אלה הסמלים שמייצגים ערכים של זמן ריצה. לדוגמה, במקרה הספציפי הזה, לכל רכיב של הפלט עם אינדקסים 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){rt0, rt1} -> (d0 - rt0, d1 - rt1),
domain:
d0 in [0, 19],
d1 in [0, 29],
rt0 in [0, 15],
rt1 in [0, 20]

שימו לב שעכשיו יש לנו rt0 ו-rt1 שמייצגים ערכים של זמן ריצה. במקרה הספציפי הזה, לכל רכיב של הפלט עם האינדקסים d0, d1, ניגשים להיסטים של הפרוסה of1 ו-of2 כדי לחשב את האינדקס של הקלט. הרווחים של משתני זמן הריצה נגזרים מתוך הנחה שכל הפלח נשאר בגבולות.

מיפוי הפלט לקלט עבור of1 ו-of2:

(d0, d1) -> (),
domain:
d0 in [0, 19],
d1 in [0, 29]

איסוף

יש תמיכה רק באיסוף הפשוט. מידע נוסף זמין בקובץ 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){rt0, rt1} -> (d1 + rt0, d2 + rt1, d3),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
rt0 in [0, 26],
rt1 in [0, 68]

שימו לב שעכשיו יש לנו סמלי rt שמייצגים ערכים של זמן ריצה.

מיפוי פלט לקלט עבור 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, *) של טנזור 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)
out = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
  dimensions={0}, to_apply=max

מיפוי הפלט לקלט:

‫out[0] -> ‏p0:

(d0)[s0] -> (s0, d0),
domain:
d0 in [0, 9],
s0 in [0, 255]

‫out[0] -> ‏p0_init:

(d0) -> (),
domain:
d0 in [0, 9]

המיפוי של הקלט לפלט:

‫p0 -> ‏out[0]:

(d0, d1) -> (d1),
domain:
d0 in [0, 255],
d1 in [0, 9]

‫p0_init -> ‏out[0]:

()[s0] -> (s0),
domain:
s0 in [0, 9]

פרוסה

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

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-D ל-1D.

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. לכן, מפות האינדוקס שלו הן רק קומפוזיציה של מפות האינדוקס לרצף הזה.

Concatenate

מיפוי הפלט לקלט של concat מוגדר לכל הקלטים, אבל עם דומיינים לא חופפים, כלומר רק אחד מהקלטים ישמש בכל פעם.

p0 = f32[2, 5, 7] parameter(0)
p1 = f32[2, 11, 7] parameter(1)
p2 = f32[2, 17, 7] parameter(2)
ROOT output = f32[2, 33, 7] concatenate(f32[2, 5, 7] p0, f32[2, 11, 7] p1, f32[2, 17, 7] p2), dimensions={1}

מיפוי הפלט לקלט:

‫output -> ‏p0:

(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 4],
d2 in [0, 6]

‫output -> ‏p1:

(d0, d1, d2) -> (d0, d1 - 5, d2),
domain:
d0 in [0, 1],
d1 in [5, 15],
d2 in [0, 6]

‫output -> ‏p2:

(d0, d1, d2) -> (d0, d1 - 16, d2),
domain:
d0 in [0, 1],
d1 in [16, 32],
d2 in [0, 6]

הקלט למפות הפלט:

‫p0 -> ‏output:

(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 4],
d2 in [0, 6]

‫p1 -> ‏output:

(d0, d1, d2) -> (d0, d1 + 5, d2),
domain:
d0 in [0, 1],
d1 in [0, 10],
d2 in [0, 6]

‫p2 -> ‏output:

(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)
output = f32[4, 128, 64] dot(p0, p1),
  lhs_batch_dims={0}, rhs_batch_dims={0},
  lhs_contracting_dims={2}, rhs_contracting_dims={1}

מיפוי הפלט לקלט:

output -> p0:

(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 -> p1:

(d0, d1, d2)[s0] -> (d0, s0, d2),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 63],
s0 in [0, 255]

הקלט למפות הפלט:

p0 -> output:

(d0, d1, d2)[s0] -> (d0, d1, s0),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 255],
s0 in [0, 63]

p1 -> 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]

Pad

הוספה לאינדקס של 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

הגדרת המרווח הפנימי 1_4_1x4_8_0 מציינת lowPad_highPad_interiorPad_dim_0 x lowPad_highPad_interiorPad_dim_1.

מיפוי הפלט לקלט:

output -> p0:

(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4),
domain:
d0 in [1, 7],
d1 in [4, 7],
(d0 - 1) mod 2 in [0, 0]

output -> p1:

(d0, d1) -> (),
domain:
d0 in [0, 11],
d1 in [0, 15]

ReduceWindow

הפעולה ReduceWindow ב-XLA מבצעת גם ריפוד. לכן, אפשר לחשב את מיפוי האינדקסים כהרכבה של אינדקס ReduceWindow שלא מבצע ריפוד, ושל אינדקס PadOp.

c_inf = f32[] constant(-inf)
p0 = f32[1024, 514] parameter(0)
outpu = f32[1024, 3] reduce-window(p0, c_inf),
  window={size=1x512 pad=0_0x0_0}, to_apply=max

מיפוי הפלט לקלט:

output -> p0:

(d0, d1)[s0] -> (d0, d1 + s0),
domain:
d0 in [0, 1023],
d1 in [0, 2],
s0 in [0, 511]

output -> c_inf:

(d0, d1) -> (),
domain:
d0 in [0, 1023],
d1 in [0, 2]

אינדוקס של מפות ל-Fusion

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

קלט אחד, כמה מפות אינדוקס

לדוגמה: 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 output = 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 מתייחס למאפיין הפנימי ביותר של הקלט.

דוגמאות נוספות זמינות בקובץ indexing_analysis_test.cc.

Indexing Map Simplifier

הכלי הפשוט לשינוי נתונים שמוגדר כברירת מחדל ב-mlir::AffineMap לא יכול להניח הנחות לגבי טווחי המאפיינים או הסמלים. לכן, הוא לא יכול לפשט ביעילות ביטויים עם mod ו-div.

אנחנו יכולים להשתמש בידע לגבי הגבולות התחתונים והעליונים של ביטויי המשנה במיפויים האפיניים כדי לפשט אותם עוד יותר.

הכלי לפישוט יכול לשכתב את הביטויים הבאים.

‫(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) למשך d ב-[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.

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

דוגמה

למה בחרנו לעשות זאת?

מיזוג

ניצול אופרנדים

פריסה

מפת אינדקס

הטמעה

אינדוקס של מפות לפעולות לא משולבות

Elementwise

‫(Variadic)Reduce

כיווץ צורה

הרחבת הצורה

שינוי צורה גנרי

דוגמה 1: ליניאריזציה – דה-ליניאריזציה.

דוגמה 2: צורות משנה מורחבות ומכווצות

Bitcast

אינדוקס של מפות ל-Fusion

קלט אחד, כמה מפות אינדוקס

קלט אחד, מפת אינדקס ללא כפילויות

Softmax

Indexing Map Simplifier

ניתוח הוספה לאינדקס