การโทรที่กำหนดเองแบบ 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 ที่ประมวลผลแบบเดียวกับโค้ด CPU ด้านบน (A[i] = B[i % 128] + C[i])

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 ที่มีพารามิเตอร์ 2 รายการ ได้แก่ ins และ out แบบแผนการเรียก GPU ทำให้ สามารถจัดการอินพุต/เอาต์พุตที่เป็นรูปร่าง Tuple ได้อย่างมีประสิทธิภาพ

อย่างในตัวอย่าง 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 และดีซีเรียลไลซ์ข้อมูลดังกล่าวภายในการเรียกใช้ 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 จะแสดง Tuple ในหน่วยความจำเป็นอาร์เรย์ของ Pointer ใน Pseudocode ของ 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

เอาต์พุต Tuple เป็นบัฟเฟอร์ชั่วคราว

การป้อนข้อมูล Tuple ไปยังการโทรที่กำหนดเองมีประโยชน์ แต่ก็ไม่ได้จำเป็นเสมอไป ถ้าเราไม่สนับสนุนอินพุต Tuple ในการเรียกที่กำหนดเอง คุณสามารถแยก tuple โดยใช้ get-tuple-element ก่อนส่งไปยังการเรียกที่กำหนดเองได้เสมอ

ในทางกลับกัน เอาต์พุต Tuple ช่วยให้คุณสามารถทำในสิ่งที่ทำไม่ได้

เหตุผลที่ชัดเจนของการมีเอาต์พุต Tuple คือเอาต์พุต Tuple คือวิธีที่การเรียกที่กำหนดเอง (หรือ XLA op อื่นๆ) แสดงผลอาร์เรย์อิสระหลายชุด

แต่แน่นอนว่าเอาต์พุตของ Tuple ก็เป็นอีกหนึ่งวิธีในการสร้างหน่วยความจำสำหรับการโทรที่กำหนดเองเช่นกัน ได้ เอาต์พุตสามารถแสดงบัฟเฟอร์ชั่วคราวได้ ลองพิจารณาบัฟเฟอร์เอาต์พุตว่ามีพร็อพเพอร์ตี้ที่ออบเจ็กต์นี้เขียนลงไปได้ และจะอ่านได้จากบัฟเฟอร์ดังกล่าวหลังจากเขียนลงไป นั่นคือสิ่งที่คุณต้องการจากบัฟเฟอร์ชั่วคราว

ในตัวอย่างด้านบน สมมติว่าเราต้องการใช้ F32[1024] เป็นบัฟเฟอร์ชั่วคราว จากนั้นเราจะเขียน HLO เหมือนกับด้านบน และเราจะไม่อ่าน tuple index 1 ของเอาต์พุตของการเรียกที่กำหนดเอง

Tuple ในการเรียกใช้ CPU ที่กำหนดเอง

เรามีฟังก์ชัน do_custom_call(const void** ins, void* out) ในโค้ด CPU ins คืออาร์เรย์ที่มีเพียงองค์ประกอบเดียว ซึ่งชี้ไปยัง param0 คุณสามารถเข้าถึงบัฟเฟอร์ย่อยของ param0 โดยการยกเลิกการอ้างอิงตัวชี้ และเข้าถึงบัฟเฟอร์ย่อยของ output_tuple ได้โดยการยกเลิกการอ้างอิง out

Tuple ในการเรียกใช้ GPU ที่กำหนดเอง

เรามีฟังก์ชัน do_custom_call(..., void** buffers, ...) ในโค้ด GPU ในกรณีนี้ buffers คือโฮสต์อาร์เรย์ของ 6 อุปกรณ์ Pointer 1 อันสำหรับบัฟเฟอร์แต่ละใบในอินพุต/เอาต์พุต ในการสร้างลิสต์แบบแฟลตลิสต์ เราจะทำซ้ำพารามิเตอร์ และเอาต์พุต และสำหรับแต่ละรายการที่เราทำการส่งผ่านแบบสั่งล่วงหน้าสำหรับรูปร่างของพารามิเตอร์ มีความชัดเจน:

// 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