נכון למועד כתיבת המאמר, StableHLO מוכן להחליף את MHLO/HLO כממשק למהדר. אפשר ליצור אותו באמצעות TensorFlow, JAX ו-PyTorch, אפשר לצרוך אותו באמצעות XLA ומספר יישומי פלאגין של PJRT של צד שלישי, והוא כולל את כל התכונות הציבוריות ש-MHLO/HLO מספקות, וגם פונקציונליות נוספת.
במסמך הזה מתוארים השלבים הבאים בפרויקט StableHLO, עם סיווג של העבודה המתמשכת שמופיעה במעקב הבעיות וסדרת מוצרים שאנחנו מתכננים לספק.
אבני הדרך הנוכחיות
נקודות הציון הנוכחיות שלנו עוקבות אחרי שתי מגמות עיקריות:
- למקסם את התועלת של StableHLO לכל קהילת OpenXLA.
- איחוד חוויית המפתחים לכל חברי OpenXLA.
- הוצאה משימוש של MHLO: ברבעון הרביעי של 2024 התחלנו לבדוק את ההוצאה משימוש הפנימית של MHLO, והעברנו שלבים שימושיים, כולל קנוניקליזציה ודפוסי תיקיות, אל StableHLO. אחרי שנוכיח שהתהליך פשוט פנימית, אנחנו מתכננים לשתף RFC עם לוחות זמנים להעברות חיצוניות ל-StableHLO. סביר להניח שהמעבר יתבצע ברבעון הראשון של 2025, ואנחנו מתכננים לתת לצוותים מספיק זמן ותמיכה כדי לעבור ל-StableHLO ברבעון הראשון של 2025.
- העברת אופטימיזציות שאינן תלויות בחומרה ל-StableHLO: בהמשך למגמה שצוינה למעלה, אנחנו רוצים ש-StableHLO יהיה המקום הטוב ביותר לאחד פישוטים של גרפים שאינם תלויים בחומרה, כדי שכל הפלאגינים של PJRT, כולל אלה שממירים מ-StableHLO ל-IR של מהדר שאינו XLA, יוכלו ליהנות מהיתרונות המקסימליים. חלק מהמטרה הזו כולל איחוד של דפוסים שמשמשים ב-Google AI Edge, בפרויקט JAX-Enzyme ובפרויקטים אחרים, כולם במאגר StableHLO. חלק מהאיחוד הזה כבר התחיל, אבל עיקר העבודה תתבצע ותסתיים ברבעון הראשון של 2025.
- פירוק לרכיבים ב-OpenXLA: התחלנו ליצור רכיבים ייעודיים ב-openxla/xla ל-HLO, שדומים להגדרה של המאגר StableHLO (מידע נוסף). כמו כן, התחלנו להעביר את כל הקצוות העורפיים של OpenXLA מאחורי הפלאגינים של PJRT. אנחנו משקיעים גם בתיקון בעיות בולטות בממשק המשתמש שאנחנו מגלים בפלאגינים האלה של PJRT, כולל דברים כמו תקשורת מדויקת לגבי גרסת StableHLO בפלאגינים של StableHLO, כדי שאפשר יהיה להשתמש בתכונות חדשות בפלאגינים חדשים באופן מיידי (מידע נוסף).
- הפעלת קומפוזיטים מקצה לקצה: ברבעון השלישי של 2024 הוספנו קומפוזיטים ל-HLO, וכך אפשרנו תמיכה מלאה בסטאק של המהדרר לצורך הפשטות. ברבעון הרביעי של 2024 לימדנו את ה-inliner של XLA על פונקציות מורכבות והוספנו מעברים ב-HLO/StableHLO להטמעה של פונקציות מורכבות לא מוכרות עם הפירוק שלהן. אנחנו בודקים עכשיו אם להוסיף ממשקי API ייעודיים של JAX ליצירת קומפוזיטים מהמסגרת (ממשקי API של PyTorch כבר קיימים), וגם להוסיף למסמכי התיעוד של Colab הנחיות לשימוש נכון בקומפוזיטים. התהליך יושלם ברבעון הרביעי של 2024.
אבני דרך קודמות
במחצית הראשונה של 2024, השקנו את StableHLO v1.0, שבו הטמענו שיפורים בעדיפות גבוהה, כולל השלמת המפרט הציבורי של קבוצת הפעולות, מתורגמן של הפניות, הרחבת התאימות לאחור ולפנים כדי לתמוך בפריסה במכשיר באמצעות בדיקת יחידות תאימות, יכולת הרחבה באמצעות פעולות מורכבות, תמיכה בתנודתיות לפי מפרט, השקת חבילת נתוני בדיקה מלאה עם תוצאות מוערכות של נתונים מוזהבים ועוד.