เอกสารนี้จะอธิบายวิธีเขียนและใช้การเรียกใช้ 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