פיתוח backend חדש עבור XLA

מדריך ראשוני זה מיועד למאמצים מוקדמים שרוצים למקד מחדש בקלות את TensorFlow לחומרה שלהם בצורה יעילה. המדריך אינו שלב אחר שלב ומניח ידע על LLVM , Bazel ו-TensorFlow.

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

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

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

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

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

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

אם אין LLVM backend קיים אבל קיים סוג אחר של מחולל קוד, זה אמור להיות אפשרי לעשות שימוש חוזר ברוב ה-CPU הקיים.

תרחיש 2: חומרה לא דמוית מעבד עם קצה אחורי של LLVM קיים

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

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

תרחיש 3: חומרה לא דמוית מעבד ללא קצה אחורי של LLVM קיים

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

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