XLA (אלגברה לינארית מואצת) הוא מהדר של למידת מכונה (ML) שמבצע אופטימיזציה של האלגברה הלינארית, ומשפר את מהירות הביצוע ואת השימוש בזיכרון. בדף הזה מופיעה סקירה כללית קצרה על היעדים והארכיטקטורה של הידור XLA.
מטרות
כיום, XLA תומך במספר ממשקי framework של למידת מכונה (כולל PyTorch, TensorFlow ו-JAX) והוא חלק מפרויקט OpenXLA – סביבה עסקית של טכנולוגיות מהדר בקוד פתוח ללמידת מכונה, שפותחה בשיתוף פעולה על ידי ארגוני חומרה ותוכנה מובילים של למידת מכונה. לפני יצירת הפרויקט OpenXLA, פותחנו את XLA בפרויקט TensorFlow, אבל היעדים הבסיסיים לא השתנו:
לשפר את מהירות הביצוע. כדאי ליצור תתי-גרפים לצמצום זמן הביצוע של פעולות לטווח קצר ולבטל פעולות תקורה מזמן הריצה, פעולות של צינורות נתיכים לצמצום תקורת הזיכרון והתמחות בצורות tensor ידועות שמאפשרות הפצה מתמדת יותר.
שיפור השימוש בזיכרון. ניתוח ותזמון של שימוש בזיכרון, ללא מאגרי אחסון זמניים רבים.
הפחתת ההסתמכות על פעולות מותאמות אישית. לבטל את הצורך בפעולות מותאמות אישית רבות על ידי שיפור הביצועים של פעולות שהתמזגו באופן אוטומטי ברמה נמוכה, כדי שיתאימו לביצועים של פעולות מותאמות אישית שמוזגו במקור באופן ידני.
לשפר את הניידות מאפשרים לכתוב בקלות קצה עורפי חדש לחומרה חדשנית, כדי שחלק גדול מהמודלים של למידת המכונה יוכלו לפעול ללא שינויים בחומרה הזו. זאת בניגוד לגישה של התמחות בפעולות מונוליתיות בודדות לחומרה חדשה, שבה נדרשת כתיבה מחדש של מודלים כדי להשתמש בפעולות האלה.
איך זה עובד
המהדר של XLA לוקח גרפים של מודלים ממסגרות של ML שמוגדרות ב-StableHLO, ויוצר מהם הוראות למכונות שמתאימות לארכיטקטורות שונות. ב-StableHLO מוגדרת קבוצת פעולות עם גרסאות (HLO = פעולות ברמה גבוהה), שמספקת שכבת ניידות בין ה-frameworks של ML לבין המהדר.
באופן כללי, תהליך ההידור שממיר את תרשים המודל לקובץ הפעלה מותאם ליעד כולל את השלבים הבאים:
XLA מבצע מספר העברות מובנות של אופטימיזציה וניתוח בתרשים ה-StableHLO שלא תלויים ביעד, למשל CSE, שילוב של פעולה בלתי תלויה ביעד וניתוח מאגר נתונים זמני להקצאת זיכרון זמן ריצה לחישוב. במהלך שלב האופטימיזציה, XLA גם ממיר את הדיאלקט היציבHLO לדיאלקט HLO פנימי.
XLA שולח את חישוב ה-HLO לקצה עורפי לצורך אופטימיזציות נוספות ברמת ה-HLO, הפעם תוך התחשבות במידע ובצרכים הספציפיים ליעד. לדוגמה, הקצה העורפי של ה-GPU עשוי לבצע מיזוגים של פעולות שיועילו במיוחד למודל התכנות של ה-GPU, ולקבוע איך לחלק את המחשוב לזרמים. בשלב הזה, הקצוות העורפיים עשויים גם להתאים לפעולות מסוימות או לשילובים מסוימים שלהן לקריאות שעברו אופטימיזציה לספרייה.
לאחר מכן, הקצה העורפי מבצע יצירת קוד לפי יעד ספציפי. הקצוות העורפיים של המעבד (CPU) וה-GPU שכלולים ב-XLA, משתמשים ב-LLVM לרמה נמוכה של IR, אופטימיזציה ויצירת קוד. הקצוות העורפיים האלה פולטים את ה-LLVM IR הנחוץ כדי לייצג את חישוב ה-HLO בצורה יעילה, ואז מפעילים את LLVM כדי לפלוט קוד נייטיב מ-LLVM IR הזה.
בתהליך הזה, המהדר (compiler) XLA הוא מודולרי בכך שקל לחבר אותו לקצה עורפי חלופי כדי לטרגט ארכיטקטורת HW חדשה. הקצה העורפי של ה-GPU תומך כרגע במעבדי GPU של NVIDIA דרך הקצה העורפי LLVM NVPTX. הקצה העורפי של המעבד (CPU) תומך במספר ISA של המעבד (CPU).