במסמך הזה מתואר המסלול של מודול XLA High Level Optimizer (HLO) מהמצב ההתחלתי שלו ועד לקובץ הפעלה סופי. לפעמים נשמיט את המילה module ופשוט נתייחס ל-HLO כאל HLO.
HLO לפני אופטימיזציה
אנחנו מתחילים עם מודול HLO שעבר אופטימיזציה מראש. ה-HLO לפני האופטימיזציה לא מכיל פעולות (ops) שנחשבות פנימיות ל-XLA, כמו fusion או bitcast. בשלב הזה, לפעולות אין פריסה, או שאם יש להן פריסה היא תתעלם ממנה. בדרך כלל, אופטימיזציה של HLO מתבצעת על ידי מסגרות ברמה גבוהה יותר, כמו TensorFlow ו-JAX. כשמשתמשים בדגל XLA -xla_dump_to, קובץ ה-HLO לפני האופטימיזציה נשמר בקובץ עם סיומת שם הקובץ before_optimizations.txt.
אופטימיזציה של מודול HLO
צינור הנתונים XLA:GPU הופך את ה-HLO שלפני האופטימיזציה ל-HLO שעבר אופטימיזציה על ידי הפעלה של רצף של מעברים. אפשר לקבץ את הכרטיסים יחד באופן סמנטי ולהריץ אותם בסדר הבא:
חלוקת כרטיסים קשורים
זה כולל כרטיסים כמו ShardyPartitioner או כרטיסים לפיצול SPMD.
שלבי אופטימיזציה
השינויים האלה יכולים לכלול גם שינויים שנועדו להפוך את התוכן לחוקי וגם שינויים שנועדו לפשט אותו.
מעברים משותפים של אופטימיזציה
בדומה למעברים לאופטימיזציה, אבל מתמקד בפעולות משותפות.
העברת הקצאת פריסה
לכל פעולת HLO מוקצה פריסה שהיא חלק מהצורה של ההוראה. הפריסה קובעת איך הטנסור מסודר פיזית בזיכרון.
דוגמה לצורה עם פריסה:
f32[10,20,30]{2,0,1}
אחרי סוג הרכיב, מופיעים הממדים הלוגיים של הצורה, ואחריהם סידור הפריסה בסדר עולה. בדוגמה הזו, המאפיין המשני ביותר הוא 30, המאפיין המשני השני הוא 10 והמאפיין הראשי הוא 20.
המטרה של הקצאת פריסות היא לצמצם את מספר ההעברות הפיזיות הנדרשות באמצעות אסטרטגיה חמדנית. הוא מתחיל עם אילוצים מסוימים של פריסה (לדוגמה, ספריות cuDNN/cuBLAS מצפות לממדים עוקבים) ומפיץ פריסות 'למטה' ואז 'למעלה' בגרף HLO. בסיום ההפצה של הפריסה, יכול להיות שחלק מההוראות יכללו פריסות סותרות, אחת שהופצה מאופרנד ואחת שהופצה ממשתמש. כדי לפתור את הבעיה, מוסיפים הוראת HLO של copy שמשנה את הפריסה מפריסת האופרנד לפריסת ההוראה.
מעברים של נירמול פריסה
מכיוון שקשה להבין את הצורה הפיזית, הפריסה
הנורמליזציה מנסה לשכתב את הצורה כך שהיא תשתמש בפריסת ברירת המחדל {rank-1, rank-2, …, 0}. בדוגמה שלמעלה, הצורה הנורמלית תהיה
f32[20,10,30]{2,1,0}. פעולות העתקה שמשנות פריסות נכתבות מחדש כשילוב של transpose ו-bitcast. מכיוון שבשלב הזה אנחנו לא יכולים לבצע נורמליזציה של כל הפעולות, יש עדיין כמה פעולות שיכול להיות שיש להן פריסות שאינן ברירת מחדל, בעיקר gather ו-dot. בגבולות בין פעולות מנורמלות לפעולות לא מנורמלות יהיו bitcast פעולות שמייצגות טרנספוזיציה, כלומר טרנספוזיציה עם פריסה שהוקצתה לה שהופכת אותה ללא פעולה מבחינה פיזית.
בנוסף, נורמליזציה של פריסות הופכת חלק מהטרנספוזיציות המרומזות לגלויות, וזה חשוב כי יצירת קוד יכולה לטפל בטרנספוזיציות גלויות באמצעות פולט ייעודי. לדוגמה, מבחינה טכנית, מותר לפעולת שינוי הצורה להיות פריסה פיזית שונה בין האופרנד לתוצאה (למשל, בגלל דרגה שונה). ההעברה ReshapeDecomposer שמופעלת כחלק מהעברות של נירמול פריסה transpose, reshape bitcast ו-transpose.
אופטימיזציה של הקצאת פריסות לאחר פרסום
הכרטיסים הכי חשובים כאן הם מיזוגים של Triton (מיזוגים של GEMM + מיזוגים של Softmax/Layernorm) או כתיבה מחדש של קריאות לספרייה. בשלב הזה מתבצע גם כוונון אוטומטי, שבו XLA בוחר בין פולטים שונים, בוחר את האלגוריתם הטוב ביותר עבור קונבולוציות או נקודות, מוצא את הריצוף הטוב ביותר עבור מיזוגים שמטופלים על ידי פולט Triton וכו'.
כרטיסי פיוז'ן
שני המעברים העיקריים הם מיזוג PriorityFusion ומיזוג Multi-Output.
ב-PriorityFusion, אנחנו יוצרים מיזוגים בהתאם למודל העלויות. כשמבצעים מיזוג, אנחנו מאפשרים לשכפל פעולות עם כמה משתמשים אם אפשר למזג את הפעולה עם כל המשתמשים. בנוסף, נאפשר הרחבה של מיזוגי Softmax קיימים של Triton אם הדבר אפשרי.
Multi-Output fusion הוא שלב נפרד שמאפשר מיזוג של פעולות או מיזוגים שמשתפים אופרנד. הוא יכול גם למזג אופרנדים או מיזוגי אופרנדים למשתמשים בלי לשכפל אותם, על ידי הוספת פלט נוסף, כך שמשתמשים אחרים באופרנד שאותו רוצים למזג יוכלו להיות מופנים לפלט הזה. המעבר הזה צריך להיזהר שלא ליצור מחזורים בגרף HLO.
אחרי מיזוג פלט מרובה, מתבצעת הסרת ביטויים משותפים (HloCSE pass)
שעשויים למזג מחדש פעולות שהוכפלו קודם לכן אם הן הסתיימו
באותו מיזוג.
כמה מעברים אחרי המיזוג
כמה כרטיסים שקשורים לקולקטיבים (למשל, הפיכתם לאסינכרוניים או אכיפת סדר יחסי מסוים של קולקטיבים).
לבסוף, אנחנו מריצים את הפקודה CopyInsertion שבה מתווספות עותקים כדי לוודא שפעולות במקום לא יגרמו להחלפת נתונים שעדיין נדרשים במקומות אחרים.
בסיום האופטימיזציה, אם משתמשים בדגל -xla_dump_to, ה-HLO שעבר אופטימיזציה נשמר בקובץ עם הסיומת "after_optimizations.txt". אם רוצים להוציא את ה-HLO אחרי מעברים ביניים שמשנים בפועל את ה-HloModule, אפשר להשתמש בדגל -xla_dump_hlo_pass_re=.* (או בביטוי רגולרי ספציפי כדי להגביל את זה למעברים מסוימים).
תזמון
למודול HLO ללא לוח זמנים עדיין יש מידה מסוימת של חופש בסדר שבו הפעולות מעובדות. כל מיון טופולוגי שמכבד את קשרי האופרנד/התוצאה ואת תלות הבקרה הוא תקף. התזמון קובע באיזה סדר ספציפי להשתמש. הדאגה העיקרית בשלב הזה היא צריכת הזיכרון המקסימלית, שתלויה במשך החיים של טנסורים. בשלב הראשוני, אנחנו מנסים אלגוריתמים שונים של מתזמן ומקצים את התזמון שאמור למזער את צריכת הזיכרון המקסימלית. שימו לב: בשלב הזה אנחנו עדיין לא עובדים עם מאגרי נתונים פיזיים (זה יקרה בשלב 'הקצאת מאגרים'), ומדמים את השימוש בזיכרון.
לאחר מכן, המערכת מריצה את LatencyHidingScheduler ומנסה למקסם את החפיפה בין החישובים לתקשורת. אבל זה עלול להגדיל שוב את השימוש בזיכרון.
לבסוף, אם צריכת הזיכרון המקסימלית גבוהה מכמות הזיכרון הזמין, אנחנו מריצים את הפקודה HloRematerialization. השלב הזה מנסה לצמצם את השימוש בזיכרון על חשבון הביצועים. לדוגמה, יכול להיות שחלק מהמיזוגים יפוצלו וחלק מהפעולות ישוכפלו כדי לקצר את משך החיים של המאגר. אם מתרחש חומרור מחדש, כדאי לבדוק דרכים לצמצום דרישות הזיכרון בצד המודל (למשל, שימוש בגדלים קטנים יותר של אצווה).
הקצאת מאגר נתונים זמני
ממש לפני שאנחנו עוברים ל-LLVM IR, אנחנו מריצים את שלבי ההקצאה של המאגרים שיקצו פרוסות של מאגרים לכל הוראה בתרשים HLO. הקצאת המאגר מתבצעת בכמה שלבים:
HloDataflowAnalysisמקצהHloValues(בעצם מאגרי נתונים לוגיים) להוראות. בפעולות במקום, אפשר לעשות שימוש חוזר ב-HloValueשל אופרנד. יכול להיות שאופרטור יגדיר יותר מ-HloValueאחד (לדוגמה, עם צורת תוצאה של tuple).
HloAliasAnalysisמנסה לשלב מאגרי נתונים זמניים לפעולות של יצירת כינויים, ומחשב מיפוי מ-HloValueל-HloBuffer.
BufferAssignmentמחשב מיפוי שלHloBuffersלפרוסות של מאגר נתונים זמני בתוך מאגר נתונים זמני גדול, כך שאותה פרוסה של מאגר נתונים זמני לא תשמש לערכים שונים שלHloBuffersעם חפיפה בפרקי הזמן שלהם. במקרים של פעולות שיכול להיות שיש להן כינוי, מותר שתהיה חפיפה קלה (זמן הסיום של פעולה אחתHloBufferיכול להיות זהה לזמן ההתחלה של פעולה אחרתHloBuffer). כשמשתמשים בדגל-xla_dump_to, חלק מהמידע על הקצאת מאגרים נשמר בקובץ עם סיומת השם after_optimizations-buffer-assignment.txt.
Thunks
אחרי שמבצעים אופטימיזציה לגרף HLO ומתזמנים אותו, הוא מועבר לרצף ליניארי של חלקי קוד (thunks) עבור קצה עורפי ספציפי (מעבד או מעבד גרפי).
ב-XLA, Thunk הוא הפשטה של יחידת עבודה עצמאית שזמן הריצה מבצע. יכול להיות שמדובר בהפעלה של ליבת הידור, בפעולה ספציפית, בקריאה לספרייה, במבנה של זרימת בקרה, בתקשורת קולקטיבית וכו'. Thunk Sequence מייצג את כל הקובץ שניתן להפעלה עבור קצה עורפי ספציפי.
פליטת Thunk
התהליך של המרת חישוב מתוזמן של HLO לרצף של thunk נקרא 'פליטת thunk'. הטיפול בזה מתבצע על ידי מחלקה ייעודית של פולט בכל קצה עורפי.
ב-GPU Backend, הפעולה הזו מתבצעת על ידי IrEmitterUnnested.
EmitHloComputation iterates through the scheduled list of HLO Instructions in
a computation and dispatches to a specialized Emit... method (e.g.,
EmitFusion, EmitConvolutionThunk, EmitWhile). כל אחת מהשיטות האלה יוצרת את אובייקט ה-Thunk המתאים ומוסיפה אותו לרצף ה-Thunk.
ב-CPU Backend, ThunkEmitter ממלא את התפקיד הזה ומאורגן באופן דומה. הסכום הסופי ThunkSequence
מוטמע ב-CpuExecutable.
שימו לב שכל הוראה בחישוב של רשומה במודול HLO עשויה להתאים לאפס (kTuple, kConstant, ..), לאחד או לכמה (לדוגמה, הוראת מיון) thunks ברצף ה-thunks הסופי.
Command Buffers: Optimizing Execution on the GPU
חומרת GPU מודרנית מאפשרת להקליט רצף של פעולות GPU (הפעלות ליבה, העתקות זיכרון וכו') פעם אחת, ואז להפעיל מחדש את הרצף מספר פעמים עם תקורה מינימלית של CPU. זהו שיפור קריטי בביצועים, במיוחד לעומסי עבודה עם הרבה ליבות קטנות שמופעלות במהירות. XLA משתמש ב-Command Buffer כהפשטה של CUDA Graphs או HIP Graphs. ממשק הליבה מוגדר ב-GpuCommandBuffer.
מאגר פקודות מיוצג ברצף thunk על ידי CommandBufferThunk.
הפונקציה ליצירת פלט לא יוצרת את ה-thunk הזה ישירות מהוראות HLO. במקום זאת, הפעולה הזו מתבצעת על ידי CommandBufferConversionPass שפועל על ThunkSequence עצמו.
המעבר מזהה רצפים רציפים של פונקציות תואמות (למשל, סדרה של KernelThunk ו-GemmThunk). לאחר מכן, הוא מחליף את רצף המשנה שנמצא ב-CommandBufferThunk יחיד. ה-thunk החדש מכיל את הלוגיקה של ה-thunks המקוריים כרשימה של אובייקטים קלי משקל מסוג CommandBufferCmd.
כש-CommandBufferThunk מופעל בפעם הראשונה בזרם נתונים נתון של GPU, הוא "מתעד" את רצף הפקודות שלו במאגר פקודות של חומרה. בכל ההרצות הבאות, היא פשוט מנפיקה פקודה אחת ל-GPU כדי לבצע 'הפעלה חוזרת' של הרצף שתועד. כך נמנעים מעומס יתר על המעבד (CPU) שנוצר מהפעלת כל ליבה בנפרד.
קובץ הפעלה
התוצר הסופי של צינור ההידור של XLA הוא קובץ הפעלה עצמאי שספציפי לפלטפורמה. האובייקט הזה מכיל את כל המידע שנדרש להרצת התוכנית שעברה קומפילציה במכשיר היעד, כמו מעבד או מעבד גרפי. הוא מהווה גשר בין הקומפיילר לבין זמן הריצה. סביבות ריצה מודרניות כמו PJRT משתמשות בהפשטות ברמה גבוהה יותר (ראו PjRtExecutable), אבל בסופו של דבר הן עוטפות קובץ הפעלה ספציפי ל-backend.
קובץ Executable מכיל כמה פרטי מידע חשובים שנוצרו במהלך ההידור. התוכן המדויק משתנה בהתאם למערכת העורפית, אבל בדרך כלל הוא כולל:
קוד שעבר קומפילציה: זהו קוד מכונה ברמה נמוכה שיפעל במכשיר. במקרה של מעבדי CPU, בדרך כלל מדובר בקובץ אובייקט אחד או יותר. ב-GPU, זהו קוד המכשיר שעבר קומפילציה בפורמט PTX או HSACO, שנטען ב-GPU בזמן הריצה.
תוכנית ביצוע (ThunkSequence): הליבה של לוגיקת זמן הריצה. זוהי רצף ליניארי של אובייקטים מסוג Thunk. כל פעולת thunk מייצגת יחידת עבודה אחת, כמו הפעלת ליבה, קריאה לפונקציית ספרייה (למשל, cuBLAS) או טיפול בזרימת בקרה. זמן הריצה מריץ את התוכנית על ידי חזרה על הרצף הזה.
פריסת הזיכרון (BufferAssignment): זהו חלק קריטי במטא-נתונים, שנוצר על ידי BufferAssigner, ומתאר את פריסת הזיכרון המלאה לחישוב. הוא מציין את הגודל של כל מאגר ואת אופן הקצאת הזיכרון והשימוש החוזר בו לפרמטרים, לפלט ולערכים זמניים. סביבת זמן הריצה משתמשת בזה כדי להקצות זיכרון למכשיר ולהעביר את המצביעים הנכונים לכל thunk.
(אופציונלי) מודול HLO: לצורך ניפוי באגים ויצירת פרופילים, קובץ ההפעלה שומר בדרך כלל הפניה ל-HloModule הסופי והמותאם שממנו הוא קומפל.
הקומפיילר מתזמן את היצירה של קובץ ההפעלה הסופי לכל קצה עורפי ספציפי. השיטה RunBackend של הטמעה של Compiler היא השלב האחרון בתהליך הקומפילציה, שבו כל הארטיפקטים המקומפלים נארזים באובייקט Executable.
GpuCompiler ו-CpuCompiler מכוונים ל-GPU ול-CPU בהתאמה.
כשמשתמש קורא ל-Execute... בקובץ הפעלה, סביבת זמן הריצה משתמשת ב-BufferAssignment כדי להקצות זיכרון, ואז מפעילה את ThunkSequence כדי להפעיל את הפעולות במכשיר באמצעות הקוד המהודר.