דף זה תורגם על ידי Cloud Translation API.
Switch to English

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

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

קריאה בהתאמה אישית על מעבד

אתה יכול ליצור הוראות HLO המייצגת שיחה בהתאמה אישית באמצעות ממשק ה- API של הלקוח של XLA. זה לא נחשף באמצעות TensorFlow נכון לכתיבה.

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

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

void do_it() {
  xla::XlaBuilder b("do_it");
  xla::XlaOp param0 =
      xla::Parameter(0, xla::ShapeUtil::CreateShape(F32, {128}), "p0");
  xla::XlaOp param1 =
      xla::Parameter(1, xla::ShapeUtil::CreateShape(F32, {2048}), "p1");
  xla::XlaOp custom_call =
      xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
                      /*output_shape=*/ShapeUtil::CreateShape(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 שונה במקצת מזו במעבד. להלן דוגמא ל- 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 = threadIdx.x * blockSize.x + gridIdx.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 block_dim = 64;
  const int64 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 שלנו do_custom_call אחראית על עבודות הגיוס ב- GPU. כאן הוא משיק גרעין CUDA, אבל הוא יכול גם לעשות משהו אחר, כמו להתקשר לקוביות.

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

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

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

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

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

מעביר טיפולים לשיחות בהתאמה אישית

שקול את השיחה המותאמת אישית הבאה.

 using xla::ShapeUtil;
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 ++ - pseudocode, פרמטר 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 הוא זהה במעבד ו- GPU, הם מטופלים באופן שונה במוסדות השיחה המותאמים אישית של ה- CPU וה- GPU.

תפוקות קלות כמאגרי זמניים

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

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

הסיבה הברורה שיש פלטים של טופל היא שכך שיחה בהתאמה אישית (או כל XLA op אחר) מחזירה מערכים עצמאיים מרובים.

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

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

טפלים בשיחות מותאמות אישית של מעבד

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

טפלים בשיחות בהתאמה אישית של 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