דף זה תורגם על ידי Cloud Translation API.
Switch to English

MLIR CodeGen ל- XLA

XLA פועלת ב- HloInstruction ומבצעת אופטימיזציות רבות לייצוג זה HloInstruction הרבה כאלה בין מכשירים ממוקדים. כנקודה מסוימת מחושב לוח זמנים לינארי ומאגר הזיכרון מוקצה לכל ערך באופן סטטי. הקודן הספציפי להתקנים פועל על ידי חציית רצף זה וקריאה "emitters" ליצירת ייצוג המתאים למכשיר (לדוגמא פונקציית LLVM יחידה לפי חישוב XLA במעבד, או רצף של "גזים" המקפיא פעולות GPU ואולי נוצר PTX כאשר מיקוד ל- GPU).

כצעד שלב, אנו נמצאים כעת בתהליך יירוט התהליך מיד לאחר ש- XLA מסיים את שלב הקצאת החיץ ופולט במקום זאת מודול lhlo בדיאלקט lhlo . משם אנו מבצעים את הקודגן באמצעות רכיבי MLIR (Linalg, affine ו- GPU-ניב בעיקר) בהתאם למכשיר.

להלן תוכנית הרשומה להעברת lhlo XLA / GPU באמצעות lhlo .

משימות

מנחה התקן
פורמט הכנסה לימוד Hlo * (משימה 1) לימוד Hlo * (משימה 1)
פורמט פלט xla :: Thunk (משימה 2) LLVM IR (משימה 3)
  • משימה 1 משנה את פורמט הזנת המארח וההתקן הן מ- HloInstruction * ל- LHLO.
  • משימה 2 משנה את פורמט הפלט של המארח מהגדלים ל"איזה משטח נחיתה למארח "(ראה להלן).
  • משימה 3 מעביר את פלט המכשיר מ- LLVM IR לצורה כלשהי של MLIR. זה אופציונלי לפרויקט זה, ועיין בפרק "העברת מכשיר LLVM IR".

פרויקט זה מתעדף את קיומם של דגמים ניתנים לקצה לקצה עם פולטות LHLO ככל האפשר. זה מרמז שרשימת ההזמנות הבאה לפי עדיפות:

  • הפוך את XLA / GPU לניצול עם פולטות LHLO, כאשר תאונקים ופולטים קיימים ללא שינוי.
  • בטל את ההפניות ל- HloInstruction * ב- LHLO, כל מקרה לגופו:
    • החלף פולט מדור קודם לפולט מבוסס MLIR (למשל לינל), או
    • תרגם מכנית את הפולט הקיים כדי לקחת ייצוג MLIR (העבר לסטנדרט באמצעות GPU Dialect).

גזים נודדים (משימה 2)

xla :: gpu :: Thunk הוא מבנה נתונים ש:

  • ניתן להזמין אותו מהמארח (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • נושא נתונים שונים בתתי המשנה שלו.
  • אינטראקציה עם BufferAllocation :: Slice ו- StreamExecutor.
  • משיק גרעינים
  • שיחות לכל ספריות זמן ההפעלה.

העלות של זה כוללת:

  • ייצוג נתוני תצורה ספציפיים (לדוגמה, קונפי קונבולציה).
  • נודדות צורות אופ וצורות אופרנד.
  • ייצוג עץ גזעים (בזמן, מצב וכו ').

עבודת ההגירה אינה תלויה בהעברת LHLO / emitter. תחת משאבים מוגבלים, העדיפות העומדת מאחורי הגירת LHLO / emitter.

יש לנו כמה אפשרויות כיצד להוריד את החלק בצד המארח מ- LHLO:

  • TFRT
    • (מקצוענים) עטיפות CUDA ו- HIP נהדרות לשימוש.
    • (מקצוען) קל ליישום שיחות ספריה (cuDNN, cuBLAS, cuFFT וכו '), כמו ש- TFRT מתפרשים על ידי קוד C ++.
    • (קון) הצד המארח נמצא בפיתוח ולא נבדק.
  • קוד מעבד מעובה
    • (מקצוען) יכולת נמוכה גדולה. צור כמה לולאות ותנאים וזה נעשה.
    • (קון) GPUDialect עדיין לא מדגם שרשראות / זרמים / אסינכרוניות / הקצאת מכשירים.
    • (Con) תמיכה בזמן ריצה של CUDA / HIP היא מינימלית (נתיב ערכת כלים, גרסה, טעינה דינמית וכו ').
  • זמן ריצה של XLA קיים

החלטה: לאמץ TFRT, אך גם לתמוך בקוד מעבד ריצה ב- TFRT.

העברת מכשיר LLVM IR (משימה 3)

פולט אלמנטים מייצר יעד אופטי על ידי מילויו של אלמנט אחר אלמנט. כל אלמנט פלט תלוי בקבוצת אלמנטים מהאופראנדים. כל האלמנטים מתוארים על ידי שילוב המאגר עם מדדים דינמיים. די לתאר כמעט את כל האופציות ה"מתמטיות ", אך מסיבות ביצועים רק קבוצת משנה גדולה של אופציות" מתמטיקה "מיושמת ישירות ב- (Cpu | Gpu) ElementalIrEmitter.

ElementalIrEmitter הוא ייחודי בכך:

  • חלק גדול מהקוד משותף בין XLA / GPU ו- CPU.
  • זה מייצג חלק גדול של אופציות שנראו בדגמים, כולל כל האופציות הנבונות באלמנטים.
  • רוב המיזוגים תלויים אך ורק ב- ElementalIrEmitter.
  • זה פשוט מבחינה מבנית, מכיוון שהוא מתאר DAG של תלות נתונים בין אלמנטים של אופ לבין רכיבי אופרנד.
  • זה בעיקר נייד ורמה גבוהה (למשל שלא כמו GPU kReduce ו- GPU kCopy).
  • תמיכת צורה דינמית קלה לפחות עבור אופציות חכמות-אלמנטים.

כעת, לכל האופציות, שנפלטות באופן בסיסי או לא, ישנם כמה טעמים של המצב הסופי של כל XLA op:

  1. קוד המכשיר נשאר כ- IR LLVM.
  2. מעבד מחדש את הפולט הישן כדי להיות כמו LHLO -> MLIR LLVM דיאלקט:
    • (עלות) תהיה עבודה זורקת אם אנו רוצים בסופו של דבר לעבור לסטנדרט.
    • (תועלת) זה קל ומכני. ניתן לעשות זאת בתקופה קצרה.
    • (תועלת) זה לא מועיל יותר בהשוואה ל (1).
  3. פולטים ישנים של Refactor כדי להיות כמו LHLO -> MLIR GPU + Standard + Loops:
    • (עלות) הרמת פולטים קיימים לתקן מציגה כמה אתגרים. יש להמיר מצביעים ו- GEPs ל- MemRefs ולתצוגות משנה. הבטחת שלמות amdgpu היא עוד אחת.
    • (עלות) XLA / GPU נשען מאוד על מטא נתונים של LLVM:
      • range עבור מדדי חסום / חוט.
      • align , dereferenceable , invariant.load , alias.scope , noalias עבור / חנויות עומס.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable עבור לולאות רצופות.
    • (תועלת) יכול להיות ארוך טווח. נייד יותר.
  4. Refactor emitters old to be LHLO -> לינאלג, ולכתוב פולטות לינאלג חדשות
    • (עלות) זה מקרה לגופו. בהשוואה לאפשרויות הקודמות, יישום חדש התואם לביצועים של XLA צריך לעבור את המידוד <-> אופטימיזציה של זרימת העבודה, שיכולה להיות עלות משמעותית עבור כמה אופציות.
    • (תועלת) ערימה אחידה; תמיכה בקהילה; הִטַלטְלוּת; פוטנציאל אופטימיזציה רב יותר.

מסקנות:

  • אל תלך על (2). (1) או (3) פשוט טובים יותר מ- (2). (2) עולה יותר מ- (1), מכיוון שהוא דורש הרבה שיקום מכני רב יותר. עם (1) אנו עדיין יכולים להשיג את המטרה לאפשר ל- XLA לאסוף פולטות MLIR. זאת על ידי ביצוע LHLO -> LLVM IR -> הפעלת פולטות מכשירים מדור קודם.
  • אפשרויות ElementalIrEmitter נועדו (4), אך לא באופן מצטבר. אין דרך לעשות זאת על ידי op, מכיוון שכל האופציות הנפלטות מהיסודות מחוברות לאותה גרף. עבודה זו יכולה לשמש גם כנקודת איחוד של מספר כוחות מתמשכים (xla / service / mlir_gpu, מחולל הגרעינים, לינאלג).
  • כל שאר האפשרויות הולכות על (1). כיעד מתיחה, הם עשויים להיות מועברים ל (3) או (4).

עדיפות

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

העדיפות היא: משימה 1 (LHLO לפולטות מדור קודם)> משימה 2 (גזים)> משימה 3 (פולטות MLIR).

בסוף המשימה 1, משתמשים ב- XLA יכולים ליצור LHLO (למשל מחולל גרעינים) ולבצע אותם. פורמט ההדרכה לא יכול להיות MLIR בהמשכים.

בסוף המשימה 2, LHLO מוריד ל- MLIR ראוי ומסדרתי. זה מאפשר אוספים לא מקוונים.

בסוף המשימה 3, כל פולטי ה- XLA מבוססים על MLIR ביישומה.

עיצוב מפורט

שלב 1: (משימה 1) השלם LHLO וגרם לפולטות מורשת לקחת LHLO

שלב זה גורם לכל פולטות XLA / GPU הקיימות לקיים אינטראקציה עם אופציות MLIR. שלב זה הוא שיקום מחדש טהור ו- NFC.

שלב זה הוא מכני ברובו, אך כדאי לשים לב לפערים הבאים בין HloComputation ללא התייחסות ל- LHLO:

  • לכל HloInstruction גישה ישירה לאופראנדים שלה (DAG של זרימת נתונים). נהפוך הוא, לכל LHLO op גישה רק למאגר האופראנד שלו (דו-חלקי בין ops למאגר). האופציות של LHLO צריכות לעבור דרך שרשראות באמצעות הגדרות שימוש כדי לגשת לאופציות האופראניות שלהן.
  • פולטי מורשת שלא נקנו, כמעט ואינם אמיתיים לגשת לאופראנדים שלהם. החריג היחיד הוא kReduce.
  • פולטות לא מורכבות מדור קודם ניגשות ל- BufferAssignment רק לקבלת פרוסות, ולא לגישה למבני נתונים של Aux כמו dataflow_analysis () או alias_analysis (). llvm_ir בונה ניתוח כינוי משלו () על בסיס מידע על פרוסות.

המסקנה היא שעל LHLO להתאים ימינה ללא טרחה משמעותית.

שלב 2: (אופציונלי) תמיכה בפרופיל

שלב זה נחוץ רק אם נתחיל להשליך חלק מההיגיון XLA Thunk (ראה השלב הבא).

לפני שנדליק למעשה פולטים מבוססי MLIR, אנו זקוקים לפרופילציה עבור פולטים מבוססי MLIR.

נכון לעכשיו XLA מבצע פרופיל משלה באמצעות התקשרות למערכת טיימר של StreamExecutor's. הטיימר מתחת למכסה המנוע מכניס שני אירועים לפני שיגור גרעינים ואחריו, ומודד את זמן הסנכרון בין שני אירועים אלה.

ישנן בערך שלוש גישות לתמיכה בפרופילציה ב- MLIR:

  • הרץ פרופיל מקצה לקצה
  • הוסף פרופיל אופ עבור כל אופציה ב- LHLO, באמצעות פרופיל מוזרק.

הגישה "מקצה לקצה" שקופה ל- MLIR, אך סובלת מאותה בעיה שגורמת ל- XLA לא להשתמש בה מלכתחילה: שיחות ספריה שנאספו על ידי פרופיל (nvprof / ...) לא יכולות בקלות להתייחס ל- HLO ops. לדוגמה, cuDNN משיקה מספר גרעינים עבור כל HLO, וקשה לדעת אילו גרעינים תואמים ל- HLO.

גישת "פרופיל המוזרק" דורשת:

  • LHLO לקחת פרופיל כפרמטר.
  • הכנס פרופיל. התחל / פרופיל. קדם לפני ואחרי כל אופציה.
  • מעבר מאותו פרופיל מוריד. {התחל, סוף} ליישום C ++.

לא ניתן לעשות את הפרופיל המדויק בקלות עבור אופציות שנוצרו על ידי MLIR, מכיוון:

  • ל- MLIR אין טיימר, וזה גם לא תלוי ב- TFRT / StreamExecutor.
  • MLIR לא מתקשר בקלות לפונקציות C עם פרמטרים מורכבים.

שלב 3: (משימה 2) הגירה של גזים

כהערה, ישנם בערך שלושה סוגים של גוזלים:

  • KernelThunk, שמשגר גרעין.
  • גשמי זרימת בקרה, שיש בהם היגיון זרימת בקרת מארח (מותנה, בזמן, עבור, רצף) והפעלת גרעיני גוף.
  • גברי ספרייה: cuDNN, cuBLAS, cuFFT, NCCL וכו '.

התוכנית היא:

  • הפוך את Thunks (de) לסידוריות.
  • עזור בשיפור TFRT למצב בו הוא יכול לתמוך בסמנטיקה זו.
  • ככל שהמדינה משתפרת, העבר את גונלי הבודדים באופן מצטבר.

פריטי פעולה אלה מוזמנים רק באופן חלקי. יש להעריך את סדר ביצוע הביצוע / ההקבלה ההנדסית בפועל.

שלב 4: (משימה 3) היגר ElementalIrEmitter

לאחר הפרופיל יהיה מוכן, אנו יכולים להשלים ולכוון את כל הפולטים המבוססים על ElementalIrEmitter ב- MLIR. לאחר מכן אנו מפעילים אותם כברירת מחדל, בהנחה שכל הפולטים המבוססים על MLIR משתמשים בזרם יחיד.

שימו לב שמועיל להעביר את ElementalIrEmitter של XLA / CPU גם מכיוון שהם חולקים חלק גדול מהקוד.

עם ביצוע כל ההשוואות והציד ביצועים (TODO: הגדרת זוגיות ביצועים), אנו מפעילים את הפולט האלמנטרי החדש המבוסס על MLIR ומוחקים את ElementalIrEmitter מדור קודם.

שלב זה מספק גם מעברי היתוך קלים (ops קינון) להגירה מאוחרת.

שלב 5: תמיכה או זרוק מרובי זרמים

איננו יכולים למחוק חלק מהפולטים עד שאנו תומכים בכך ב- MLIR, או נפיל את התכונה. זו כמות עבודה יחסית גדולה ב- MLIR וסכום רווח קטן עבור XLA. עלינו לחקור משתמשים נוכחיים של משתמשי XLA / GPU רב-זרם, ולנסות למחוק תכונה זו אם היא סבירה.

שלב 6: (משימה 3) אופטי מכשירים שהועברו

שלב זה מעביר את כל האפשרויות שלא נבדקו, ואז נוכל למחוק את כל הפולטים שאינם מעוניינים.

זה קורא לשכתב / מחזיר מחדש עבור kCopy ו- kReduce. kReduce כבר עובד על הרבה, כך שנשאר לראות את כמות העבודה האמיתית שצריך לעשות.