הוספת כינויים בפורמט XLA

במסמך הזה נתאר את ה-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.

כדי לציין את הכינויים באופן פרוגרמטי, ראו את XlaBuilder::SetUpAlias API.

הגדרת כינוי בזמן ריצה

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

מאגרי הנתונים הזמניים של התוכנה קשורים ExecutionInput שנ', שמכיל עץ של MaybeOwningDeviceMemory. אם הזיכרון שצוין כבעלות (הבעלות על מאגר הנתונים הזמני מועברת אל סביבת זמן הריצה של XLA), מאגר הנתונים הזמני למעשה תורם, והעדכון מתבצע, כפי בקשה שנשלחת ל-API של יצירת אימייל חלופי בזמן הידור.

עם זאת, אם מאגר הנתונים הזמני שמשתמשים בו בכינוי בזמן הידור לא נתרם במסגרת זמן ריצה, הגנה מפני העתקה מתחילה: מוקצה מאגר פלט נוסף O, והתוכן של מאגר הנתונים הזמני של הקלט P שנועד ליצור כינוי מועתק אל O (בדיוק כך שהתוכנית יכולה לפעול כאילו מאגר הנתונים הזמני O) נתרמו בזמן הריצה).