במאמר הזה מתואר ה-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
s, וכתוצאה מכך מכילים עץ MaybeOwningDeviceMemory
. אם הזיכרון מוגדר כבעלות (הבעלות על מאגר הנתונים הזמני מועברת אל זמן הריצה של XLA), המאגר נתרם בפועל והעדכון מתבצע במקומו, בהתאם לדרישה של ממשק ה-API של כינוי בזמן ההידור.
עם זאת, אם המאגר שנקרא בזמן הידור לא נתרם בזמן הריצה, הפונקציה copy-Protection מופעלת: יוקצה עוד שטח אחסון זמני לפלט O
, והתוכן של מאגר הנתונים הזמני P
שהיה אמור לשמש לכינוי מועתק אל O
(כך שהתוכנית תוכל לפעול כאילו נתרמה המאגר O
בזמן הריצה).