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

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

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

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

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

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

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

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

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

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

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

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