עיצוב לשירותי תרגום

מודל נתונים

תוכנות StableHLO הן חישובים מעל טנסורים (מערכים של n-ממדים), שבמודל הנוכחי, מיושמים באמצעות מחלקה Tensor. סוג האחסון (storage class) הבסיסי של אובייקט Tensor, detail::Buffer, שומר את mlir::ShapedType של הטנזור יחד עם אובייקט mlir::HeapAsmResourceBlob שמייצג blob ניתן לשינוי של טנזור הנתונים פרוסים כמערך בייטים רציפים סדר ראשי לקטן. האובייקטים מסוג detail::Buffer נספרים בהפניות כדי לפשט את ניהול הזיכרון.

רכיבים בודדים של טנזור מיוצגים באמצעות המחלקה Element, משתמש באיחוד מובחן שמחזיק באחד מהערכים APInt, APFloat או pair<APFloat,APFloat> לאחסון. הפריט האחרון משמש לאחסון רכיבים וגם סוגים מורכבים.

ל-Tensor יש את ממשקי ה-API הבאים כדי לקיים אינטראקציה עם הרכיבים הנפרדים שלו:

• Element Tensor::get(llvm::ArrayRef<int64_t> index): כדי לחלץ רכיב רכיב img_tensor בודד באינדקס רב-מימדי index בתור Element לאובייקט.
• void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: כדי לעדכן אובייקט Element element לארגומנט רב-מימדי להוסיף את index לאינדקס.

איך מתרגמים מבצעים

פונקציית הכניסה למפענח היא

SmallVector<Tensor> eval(func::FuncOp func, ArrayRef<Tensor> args);

שמבצעת את הפעולות הבאות:

1. מעקב אחרי ארגומנטים של SSA של func ושל זמן הריצה המשויך להם Tensor ערכים, שצוינו ב-args, באמצעות מפת טבלת סימנים, M.
2. לכל פעולה בתוך func, בסדר SSACFG:
   • הפעלה של eval בפעולה לכל אופרנד SSA של הפעולה, מחלצים את ערך של זמן ריצה מ-M שיסופק כארגומנט להפעלה של eval.
   • עוקב אחר תוצאות ה-SSA של הפעולה והערך המוערך ב-M.

רמת התפעול eval שצוינה ב-(2) אחראית ליישום של הסמנטיקה של ביצוע בהמשך מוצגת דוגמה של stablehlo::AddOp. בדוגמה, רכיבים בודדים של המעבדים lhs ו-rhs מוצגים בזוג חולץ כ-Element אובייקטים ולאחר מכן נוספים. התוצאה של ההוספה, אובייקט Element, מאוחסן בארגומנט result הסופי.

Tensor eval(AddOp op, const Tensor &lhs, const Tensor &rhs) {
  Tensor result(op.getType());

  for (auto it = result.index_begin(); it != result.index_end(); ++it)
    result.set(*it, lhs.get(*it) + rhs.get(*it));

  return result;
}

באופן כללי, העיצוב של המתורגמן עבר אופטימיזציה לשיפור הקריאוּת של פונקציות eval בפעולות נפרדות, כי הוא אמור משמשת כיישום עזר של StableHLO. לדוגמה, במקום הגדרת eval כפונקציה של תבנית והתאמתה לפרמטרים באמצעות סוגי רכיבים, כוללים פרטים על אופן הטיפול בסוגים שונים של רכיבים Element::operator+ וכו', כדי לפשט את ההטמעה של eval

שימוש בתרגום לקיפול קבוע

אנחנו יכולים להשתמש במנגנון התרגום כדי לקפל פעולות באמצעות אופרנד קבוע ערכים. קטע הקוד הבא מדגים רעיון להטמעה לקיפול של stablehlo::AddOp עם אופרנדים מסוג נקודה צפה (floating-point):

OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
  auto attrs = adaptor.getOperands();
  DenseElementsAttr lhsData = dyn_cast<DenseElementsAttr>(attrs[0]);
  DenseElementsAttr rhsData = dyn_cast<DenseElementsAttr>(attrs[1]);
  if (!lhsData || !rhsData) return {};

  auto lhs = Tensor(lhsData);
  auto rhs = Tensor(rhsData);
  auto result = eval(*this, lhs, rhs);

  SmallVector<APFloat> values;
  for (auto i = 0; i < result.getNumElements(); ++i) {
    Element element = result.get(i);
    values.push_back(cast<FloatAttr>(element.getValue()).getValue());
  }

  return DenseElementsAttr::get(result.getType(), values);
}

בשלב זה אנחנו לא עובדים באופן פעיל על שילוב של המתרגם מתקפלות כל הזמן כי אנחנו לא מתכננים להטמיע תיקייה עבור StableHLO. עם זאת, בעתיד, אנחנו מתכננים להשתמש בתרגום השיחה באופן קבוע מתקפלות ב-MHLO, ובשלב זה נשפר את הארכיטקטורה של קטע הקוד שלמעלה (למשל, יכולה להיות לנו פונקציה מסייעת שמשלבת אופרנדים קבועים בתוך Tensor אובייקטים ופרוק את Tensor תוצאות לתוך OpFoldResult).

בדיקת המתרגם של StableHLO

רכיב התרגום לוקח כקלט (A) תוכנית StableHLO ו-(B) ערכי נתונים יוזם בתוכנית, ותיצור ערכים של נתוני פלט, שמותאמים לעומת ערכי הנתונים הצפויים שהמשתמשים סיפקו. ערכי הנתונים (B) הם לבצע את הקוד בתוך התוכנה עצמה באמצעות פעולות של stablehlo.constant. המתרגם מעריך את תוכנית הקלט. הפלט של הפעולה בבדיקה נבדק באמצעות בדיקות (למשל check.expect_eq, check.expect_almost_eq), וכן שמוצגת בהמשך. check.expect_eq ו-check.expect_eq_const בודקים ברמת הסיביות שוויון לכל סוג נתמך וגם check.expect_almost_eq check.expect_almost_eq_const בודק אם יש קרוב לשוויון במסגרת סובלנות, מוסבר בהנחיה (G6) לבדיקה של נקודה צפה (floating-point) וסוגים מורכבים.

// CHECK-LABEL: Evaluated results of function: add_op_test_ui4
func.func @add_op_test_ui4() {
  %0 = stablehlo.constant dense<[0, 2]> : tensor<2xui4>
  %1 = stablehlo.constant dense<[15, 3]> : tensor<2xui4>
  %2 = stablehlo.add %0, %1 : tensor<2xui4>
  check.expect_eq_const %2, [15, 5] : tensor<2xui4>
  func.return
}

כלי בדיקה stablehlo-translate --interpret (קוד) אחראי לניתוח התוכנית, ולפרש כל פונקציה, כולל הפעולות שמרכיבים את הפונקציה. יש לנו קבוצה ייעודית לבדיקות, של מספר בדיקות שביצעו התנהגויות שונות בזמן ריצה, בכל פעולה של StableHLO. הבדיקות זמינות כאן.

הנחיות לבדיקה

(G1) האם צריך לבדוק את כל הסוגים הנתמכים בכל פעולה?

כדי להחליט, נוכל להשתמש בשילוב של הכללים הבאים:

1. במהלך הטמעה של פעולה, אם קיים קוד ב-eval המתאים לטיפול בסוג מסוים, חייבים לבצע בדיקות בנושא מהסוג הזה. לדוגמה, באופ add יש קוד בלעדי לטיפול בסוגים שלמים, בוליאני, נקודה צפה (floating-point) ומורכבים, צריך בדיקה אחת לכל קטגוריה של סוגים.

2. אם קבוצה של סוגים מטופלת בצורה אחידה בפונקציה eval המתאימה, בדיקה אחת לכל הסוגים האלה אמורה להספיק. לדוגמה, לפעולה add, כל הווריאציות של סוגי מספרים שלמים (si4, u4, si8, u8 וכן הלאה) מטופלים באותו זמן באמצעות ממשקי API של llvm::APInt, ולכן אנחנו יכולים לדלג להוסיף בדיקות לכל אחת מהווריאציות, ובמקום זאת להוסיף בדיקה בדיקת נציג. כדי למנוע אי-בהירות בבחירת הנציג, אנחנו צריך להשתמש בהנחיות הבאות:

   • אם כל הסוגים מטופלים באופן אחיד, יש להם אותו סוג פרימיטיבי (כלומר, אם כל הסוגים הם מספר שלם, נקודה צפה (floating-point) או סוגים מורכבים), לבחור את הקוד עם רוחב הסיביות המקסימלי.
   • אם כל הסוגים, בטיפול אחיד, יש שילוב של סוגים פרימיטיביים, אז לבחור את סוג המפענח הבא, בסדר יורד של preferences: integer, float-point, בוליאני, מורכב.

(G2) איך אנחנו קובעים את מספר הבדיקות הנחוצות כדי לכסות של התנהגות המשתמשים?

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

(G3) מה דעתך להוסיף בדיקות לתשתית של התכונה 'תרגום שיחה פעילה'?

התשתית של המתורגמן היא בדרך כלל פשוטה יותר ואפשר להוסיף אותה בסיס האמון שלנו. החלק הלא טריוויאלי היחיד הוא האופן שבו סוגים שונים נדחסים ונפרדים מהאחסון של המתורגמן הבסיסי. כמו שנאמר בפנייה (G1), יבדוק רק את סוגי התפעול שמטופלים באופן שונה. ב- ייתכן שקוד האריזה/האריזה, שתואם משתנים של מספרים שלמים/נקודות צפה (floating-point), ייתכן שלא יהיו מכוסים באופן מלא בדיקה. כדי להבטיח כיסוי מלא, אנחנו יכולים לבחור פעולה כמו constant תמיכה בכל סוגי הרכיבים של StableHLO וכתיבת בדיקות מקיפות.

(G4) אם ההטמעה של פעולה תלויה בפעולות אחרות, האם צריך לכתוב כמה בדיקות אפשר לעשות?

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

(G5) האם צריך לכתוב בדיקות כדי לממש את ההטמעה הקיימת / לא מוגדרת התנהגויות?

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

(G6) בזמן כתיבת בדיקות לסוגים של נקודות צפה (floating-point), רמת הדיוק צריך לציין את התוצאה הצפויה בבדיקות?

לפעולות בסיסיות (חיבור, חיסור, כפל, חילוק ריבוע), יישום בהתאם למפרט IEEE צפוי לספק התוצאה המעוגלת בטווח של 0.5 ULP מהתוצאה המתמטית המדויקת. עם זאת, לדמיין בבטחה את התוצאה הצפויה מהפעולות האלה הכי הרבה ULP אחד בנפרד. אבל יכול להיות שהשיטה הזו לא תפעל בפונקציות טרנסצנדנטליות (sine, cosine וכו') שבהן מובטחות הדיוק מוגדר על ידי ההטמעה (הסיבה).

היישום הנוכחי משתמש בשיטת "one-size-fits-all" שערך הסבלנות 0.0001. הדוגמה הבאה ממחישה את הסבלנות שצוינה למעלה בפעולה.

func.func @check_tolerance() {
  %0 = stablehlo.constant dense<0.2> : tensor<f32>

  // The following check succeeds as %0 is almost equal to the provided
  // constant modulo the tolerance, mentioned above.
  check.expect_almost_eq_const %0, dense<0.19999> : tensor<f32>

  // The following check fails as %0 is not bitwise equal to the provided
  // constant.
  check.expect_eq_const %0, dense<0.19999> : tensor<f32>

  func.return
}

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

(G7) יש לך משהו בנוגע לסגנון התכנות של הבדיקות?

1. חשוב להשתמש בשם האמיתי של הקלט/פלט במקום בברירת המחדל לערכי SSA (למשל %0, %1 וכו')
2. חשוב לוודא שהבדיקות משתמשות בפורמט הדפסה יפה, אם יש כזה.

(G8) האם צריך לכלול את הדוגמה שכבר סופקה במפרט? כן (להשלמת הבדיקה).