פיתוח קצה עורפי חדש ל- XLA

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

XLA מספק ממשק מופשט שארכיטקטורה או מאיץ חדשים יכולים להטמיע כדי ליצור קצה עורפי שיריץ פלט של תוכניות למידת מכונה באמצעות XLA. הטירגוט מחדש XLA אמור להיות פשוט וקל יותר להתאמה מאשר הטמעה של כל פעולה קיימת מ-framework של ממשק קצה כמו PyTorch או TensorFlow לחומרה חדשה.

רוב ההטמעות יתרחשו באחד מהתרחישים הבאים:

  1. הארכיטקטורה של המעבד (CPU) הקיימת עדיין לא נתמכת באופן רשמי על ידי XLA, עם או בלי קצה עורפי קיים של LLVM.
  2. חומרה שאינה דמוית-מעבד (CPU) עם קצה עורפי של LLVM.
  3. חומרה שאינה דמוית-מעבד (CPU) ללא קצה עורפי של LLVM.

תרחיש 1: ארכיטקטורה קיימת של מעבד (CPU) עדיין לא נתמכת באופן רשמי על ידי XLA

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

אם לספק החומרה יש קצה עורפי של LLVM לחומרה, קל לקשר את הקצה העורפי ל-LLVM שנוצר באמצעות XLA. במצב JIT, הקצה העורפי XLA CPU מייצר קוד של המעבד (CPU) המארח. כדי לבצע הידור מראש, הקוד xla::AotCompilationOptions יכול לספק שלשה של LLVM להגדרה של ארכיטקטורת היעד.

אם אין קצה עורפי של LLVM אבל קיים סוג אחר של מחולל קוד, אמורה להיות אפשרות להשתמש שוב ברוב הקצה העורפי של המעבד (CPU).

תרחיש 2: חומרה שאינה דמוית-מעבד (CPU) עם קצה עורפי קיים של LLVM

אפשר ליצור מודלים של הטמעה חדשה של xla::Compiler במחלקות הקיימות של xla::CPUCompiler ו-xla::GPUCompiler, כי הן כבר פולטות LLVM IR. בהתאם לאופי החומרה, ייתכן שיהיה צורך לשנות היבטים רבים של יצירת LLVM IR, אבל הרבה קוד יכול להיות משותף עם הקצוות העורפיים הקיימים.

דוגמה טובה היא הקצה העורפי של ה-GPU XLA. הקצה העורפי של ה-GPU מטרגט ISA שלא דומה למעבד (CPU), ולכן היבטים מסוימים של יצירת הקוד שלו ייחודיים לדומיין של ה-GPU. סוגי חומרה אחרים, כמו פלטפורמות DSP כמו משושה (עם קצה עורפי של LLVM במעלה), יכולים לעשות שימוש חוזר בחלקים של לוגיקת הפליטה של LLVM IR, אבל חלקים אחרים יהיו ייחודיים.

תרחיש 3: חומרה שאינה דמוית-מעבד (CPU) ללא קצה עורפי קיים של LLVM

אם אי אפשר להשתמש ב-LLVM, האפשרות הטובה ביותר היא להטמיע קצה עורפי חדש עבור XLA לחומרה הרצויה. באפשרות הזו נדרש מאמץ רב ביותר. אלה המחלקות שצריך להטמיע:

  • StreamExecutor: במכשירים רבים, לא כל השיטות של StreamExecutor נחוצות. לפרטים נוספים, אפשר לעיין בהטמעות הקיימות של StreamExecutor.
  • xla::Compiler: המחלקה הזו כוללת את ההידור של חישוב HLO ל-xla::Executable.
  • xla::Executable: המחלקה הזו משמשת להפעלת חישוב הידור בפלטפורמה.
  • xla::TransferManager: המחלקה הזו מאפשרת לקצוות העורפיים לספק מנגנונים ספציפיים לפלטפורמה לבניית נתונים ליטרליים של XLA מנקודות אחיזה נתונים בזיכרון המכשיר. במילים אחרות, הוא עוזר לבצע אנקפסולציה של העברת הנתונים מהמארח אל המכשיר ובחזרה.