במאמר הזה מתואר ה-API של כינוי XLA, שמאפשר לציין את הקישור בין מאגר הקלט למאגרי הפלט בזמן הפיתוח של תוכנת XLA.
הגדרת כינוי בזמן הידור
לדוגמה, נניח שיש מודול HLO טריוויאלי שמוסיף את 1 לקלט שלו:
HloModule increment
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
המודול הזה יקצה שני מאגרים של 4 בייטים: אחד לקלט %p ואחד לפלט %out.
עם זאת, לעיתים קרובות רצוי לבצע את העדכון במקום (לדוגמה, אם בממשק הקצה שיוצר את הביטוי, משתנה הקלט כבר לא קיים אחרי החישוב, כמו בהגדלה p++).
כדי לבצע עדכון כזה ביעילות, אפשר לציין את כינוי הקלט:
HloModule increment, input_output_alias={ {}: 0 }
ENTRY entry {
%p = f32[] parameter(0)
%c = f32[] constant(1)
ROOT %out = f32[] add(%p, %c)
}
הפורמט מציין שלפלט כולו (מסומן ב-{}) יש כינוי לפרמטר הקלט 0.
על מנת להגדיר את הכינוי באופן פרוגרמטי, עליכם לעיין ב-API של XlaBuilder::SetUpAlias.
הגדרת כינוי בזמן ריצה
הכינוי שהוגדר בשלב הקודם מצוין במהלך ההידור.
במהלך הביצוע, תוכלו להשתמש ב-API של LocalClient::RunAsync כדי לבחור אם לתרום את המאגר.
לתוכנית יש מאגרי קלט שתחומים ב-ExecutionInput, והם מכילים עץ של MaybeOwningDeviceMemory. אם הזיכרון מצוין בתור בעלות (הבעלות על המאגר מועברת לזמן הריצה של XLA), המאגר בפועל נתרם והעדכון מתבצע בהתאם לדרישה של ה-API של כינוי בזמן ההידור.
עם זאת, אם המאגר הזמני המוגדר ככינוי בזמן ההידור לא נתרם בזמן הריצה, מתבצעת הפעלה של הגנה על העתקה: מוקצית תוספת למאגר הפלט O, והתוכן של מאגר הקלט P שהיה אמור לשמש ככינוי מועתק אל O (כך שהתוכנית תוכל לפעול כאילו המאגר O נתרם בזמן הריצה).