עיצוב תרגום

מודל נתונים

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

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

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

• Element Tensor::get(llvm::ArrayRef<int64_t> index): כדי לחלץ רכיב tensor יחיד באינדקס הרב-ממדי index כאובייקט Element.
• void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: כדי לעדכן את אובייקט Element element ל-tensor באינדקס הרב-מימדי 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. בדוגמה, מתבצע חילוץ של רכיבים נפרדים של ה-Tenors lhs וה-rhs כ-Element אובייקטים, שאחר כך מוסיפים אותם. התוצאה של ההוספה, אובייקט Element, מאוחסנת ב-Tenor 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 בתפעול אישי, כי הוא אמור לשמש כהטמעת הפניה ל-SableHLO. לדוגמה, במקום להגדיר את eval כפונקציית תבנית ולהגדיר לו פרמטרים עם סוגי רכיבים, אנחנו מבצעים אנקפסולציה של פרטים על אופן הטיפול בסוגי רכיבים שונים ב-Element::operator+ וכו', כדי לפשט את ההטמעה של eval.

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

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

OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
  auto attrs = adaptor.getOperands();
  DenseElementsAttr lhsData = attrs[0].dyn_cast<DenseElementsAttr>();
  DenseElementsAttr rhsData = attrs[1].dyn_cast<DenseElementsAttr>();
  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(element.getValue().cast<FloatAttr>().getValue());
  }

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

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

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

המתורגמן לוקח את ערכי הקלט (א) תוכנית StableHLO, ו-(ב) ערכי נתונים כדי להזין לתוכנית, ויוצר ערכים של נתוני פלט, שמתבצעת בהם התאמה לערכי הנתונים הצפויים שהמשתמשים סיפקו. ערכי הנתונים (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), לנקודה צפה ולסוגים מורכבים.

// 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 Op. אפשר למצוא את הבדיקות כאן (למשל interpret_*.mlir).

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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) האם לכלול את הדוגמה שכבר סופקה במפרט? כן (להשלמת הבדיקה).