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