שיחות בהתאמה אישית XLA

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

יצירת שיחה מותאמת אישית במעבד (CPU)

אפשר ליצור הוראת HLO שמייצגת הפעלה מותאמת אישית דרך ממשק ה-API של הלקוח של XLA. לדוגמה, הקוד הבא משתמש בקריאה בהתאמה אישית כדי לחשב את A[i] = B[i % 128]+ C[i] במעבד (CPU). (ברור שאתם יכולים – ואתם רוצים! – עשו זאת עם HLO רגיל.)

#include "xla/client/xla_builder.h"
#include "xla/service/custom_call_target_registry.h"

void do_it() {
  xla::XlaBuilder b("do_it");
  xla::XlaOp param0 =
      xla::Parameter(&b, 0, xla::ShapeUtil::MakeShape(xla::F32, {128}), "p0");
  xla::XlaOp param1 =
      xla::Parameter(&b, 1, xla::ShapeUtil::MakeShape(xla::F32, {2048}), "p1");
  xla::XlaOp custom_call =
      xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                      /*shape=*/xla::ShapeUtil::MakeShape(xla::F32, {2048}));
}

void do_custom_call(void* out, const void** in) {
  float* out_buf = reinterpret_cast<float*>(out);
  const float* in0 = reinterpret_cast<const float*>(in[0]);
  const float* in1 = reinterpret_cast<const float*>(in[1]);
  for (int i = 0; i < 2048; ++i) {
    out_buf[i] = in0[i % 128] + in1[i];
  }
}
XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "Host");

שימו לב שהפונקציה do_custom_call צריכה לדעת את הממדים של החוצצים שבהם היא פועלת. בדוגמה הזו אנחנו מזינים בתוך הקוד את הגדלים 128 ו-2048. אם לא רוצים לעשות זאת, אפשר להעביר את המאפיינים כפרמטרים להפעלה.

יצירת שיחה מותאמת אישית ב-GPU

מסגרת הקריאה המותאמת אישית של ה-GPU שונה קצת מזו שבמעבד (CPU). לפניכם דוגמה ל-CUDA שעושה את אותו החישוב (A[i] = B[i % 128] + C[i]) כמו הקוד של המעבד (CPU) שמופיע למעלה.

void do_it() { /* same implementation as above */ }

__global__ custom_call_kernel(const float* in0, const float* in1, float* out) {
  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
  out[idx] = in0[idx % 128] + in1[idx];
}

void do_custom_call(CUstream stream, void** buffers,
                    const char* opaque, size_t opaque_len) {
  const float* in0 = reinterpret_cast<const float*>(buffers[0]);
  const float* in1 = reinterpret_cast<const float*>(buffers[1]);
  float* out = reinterpret_cast<float*>(buffers[2]);

  const int64_t block_dim = 64;
  const int64_t grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim,
                       /*dynamic_shared_mem_bytes=*/0, stream>>>(in0, in1, out);
}
XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "CUDA");

קודם כול שימו לב שפונקציית הקריאה המותאמת אישית של ה-GPU היא עדיין פונקציה שמופעלת במעבד (CPU). הפונקציה של המעבד (CPU) do_custom_call אחראית להוספת העבודה על ה-GPU לתור. כאן היא מפעילה ליבת CUDA, אבל היא יכולה גם לעשות משהו אחר, כמו לקרוא ל-cuBLAS.

buffers הוא מערך של מצביעים שנמצאים על המארח וכל רכיב שהוא מכיל נקודות לזיכרון המכשיר (כלומר GPU). הפרמטרים מופיעים קודם ואחריהם ערך הפלט. התהליך הזה שונה מאוד ומוסכמת הקריאה למעבד (CPU), הכוללת שני פרמטרים: ins ו-out. מוסכמת הקריאה ל-GPU מאפשרת לטפל ביעילות בקלט/פלט בצורת תלת מימד.

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

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

std::string opaque = "...";
xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                /*output_shape=*/xla::ShapeUtil::MakeShape(xla::F32, {2048}),
                opaque);

מכיוון שיש ב-xla::Shape ייצוג של מאגר נתונים זמני, אפשר לאחסן את האב האחיד שעבר תהליך סריאלי בתוך opaque ולבצע פעולת deserial שלו במסגרת הקריאה המותאמת אישית של ה-GPU. עם זאת, שימו לב שלמרות ש-xla::ShapeProto לא משתנה לעיתים קרובות, הוא כן משתנה. תוכלו לעיין ביומן ה-Git כדי לראות איך הוא השתנה בעבר.

סימון שגיאה

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

במעבד (CPU):

#include "xla/service/custom_call_status.h"

void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status);

ב-GPU:

#include "xla/service/custom_call_status.h"

void do_custom_call(CUstream stream, void** buffers, const char* opaque,
                    size_t opaque_len, xla::XlaCustomCallStatus* status);

אפשר לסמן כישלון באמצעות XlaCustomCallStatusSetFailure, למשל:

void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status) {
  // ... do some work.

  if (bad_condition) {
    char* error_message = "An error occurred";
    XlaCustomCallStatusSetFailure(status, error_message, strlen(error_message));
    return;
  }

  // ... continue.
}

אפשר גם להשתמש ב-XlaCustomCallStatusSetSuccess כדי לציין הצלחה, אבל הפרמטר XlaCustomCallStatus נמצא במצב הצלחה כברירת מחדל, כך שהתעלמות ממנו לחלוטין מציינת גם הצלחה.

כשמשתמשים בפונקציות קריאה בהתאמה אישית עם החתימה הזו, צריך ליצור את הפעולה custom-call התואמת עם גרסת ה-API המתאימה. למשל:

xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                /*output_shape=*/xla::ShapeUtil::MakeShape(F32, {2048}),
                opaque, /*has_side_effect=*/false,
                /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
                /*schedule=*/xla::CustomCallSchedule::SCHEDULE_NONE,
                /*api_version=*/API_VERSION_STATUS_RETURNING);

במקרה של כשל, לא ייעשה שימוש בפלט בהתאמה אישית של שיחות. זמן הריצה של XLA יגרום לסיום החישוב. חישוב HLO לא יכול להתאושש מהשגיאה (למשל על ידי איתור וטיפול בה).

העברת צמדים לשיחות בהתאמה אישית

כדאי לזכור את הקריאה המותאמת אישית הבאה.

using xla::ShapeUtil;
using xla::F32;
Shape p0_shape = ShapeUtil::MakeTuple({
    ShapeUtil::MakeShape(F32, {32}),
    ShapeUtil::MakeTuple({
        ShapeUtil::MakeShape(F32, {64}),
        ShapeUtil::MakeShape(F32, {128}),
    }),
    ShapeUtil::MakeShape(F32, {256}),
});
xla::XlaOp p0 = xla::Parameter(0, p0_shape, "p0");

Shape out_shape = ShapeUtil::MakeTuple({
  ShapeUtil::MakeShape(F32, {512}),
  ShapeUtil::MakeShape(F32, {1024}),
});
xla::CustomCall(&b, "do_custom_call", /*operands=*/{p0}, out_shape);

גם במעבד (CPU) וגם ב-GPU, צמד מיוצג בזיכרון כמערך של מצביעים. בפסאודו-קוד C++, פרמטר 0 שלמעלה מוצג באופן הבא.

// In-memory layout of parameter 0 from custom call above. True on both CPU
// and GPU.
float* subbuf0 = new float[32];
float* subbuf1 = new float[64];
float* subbuf2 = new float[128]
float* subbuf3 = new float[256];

void* subtuple = new void*[2];
(*subtuple)[0] = subbuf1;
(*subtuple)[1] = subbuf2;

void* p0 = new void*[3];
(*p0)[0] = subbuf0;
(*p0)[1] = subtuple;
(*p0)[2] = subbuf3;

על אף שהייצוג בזיכרון של tuples זהה ב-CPU וב-GPU, השימוש בהם נעשה באופן שונה מוסכמות הקריאה בהתאמה אישית של המעבד (CPU) ושל ה-GPU.

הפקת פלטים ב-Ttuple כמאגרי אחסון זמניים

אתם יכולים להוסיף בקלות את מערכי הנתונים הכפולים לשיחות בהתאמה אישית, אבל לא חובה לבצע אותם. אם לא תמכו בקלט כפול לקריאות מותאמות אישית, תמיד תוכלו לפתוח את החיבורים באמצעות רכיב get-tuple-tuple לפני העברתם לקריאה המותאמת אישית.

מצד שני, פלטים כפולים מאפשרים לבצע פעולות שלא ניתן לבצע אחרת.

הסיבה הברורה לפלטים כפולים היא שפלטים ב-tuple הם האופן שבו קריאה בהתאמה אישית (או כל פעולה אחרת XLA) מחזירה מערכים בלתי תלויים מרובים.

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

בדוגמה שלמעלה, נניח שרצינו להשתמש ב-F32[1024] כמאגר זמני. לאחר מכן נכתוב את ה-HLO בדיוק כמו למעלה, ופשוט לא נקרא את אינדקס Tuple 1 של פלט הקריאה המותאמת אישית.

שימוש ב-TCPU בשיחות בהתאמה אישית מהמעבד (CPU)

בקוד של המעבד (CPU), יש לנו את הפונקציה do_custom_call(const void** ins, void* out). ins הוא מערך שמכיל רק רכיב אחד, שמפנה אל param0. אפשר לגשת למאגרי המשנה של param0 על ידי ביטול האזכור של הסמן הזה. כדי לגשת למחנקים המשניים של output_tuple, צריך לבטל את האזכור של out.

Tuples בקריאות מותאמות אישית ל-GPU

בקוד ה-GPU יש לנו את הפונקציה do_custom_call(..., void** buffers, ...). במקרה הזה, buffers הוא מערך מארח של שישה מצביעי מכשירים, אחד לכל מאגר של עלים בקלט/פלט. כדי ליצור רשימה שטוחה, אנחנו עוברים על הפרמטרים והפלט, ועבור כל אחד מהם אנחנו מבצעים מעבר הזמנה מראש של הצורה. באופן מדויק:

// Layout of `buffers` parameter to GPU custom call function for custom-call
// above.
buffers[0] == subbuf0
buffers[1] == subbuf1
buffers[2] == subbuf2
buffers[3] == subbuf3
buffers[4] == output_subbuf0
buffers[5] == output_subbuf1