หน้านี้ได้รับการแปลโดย Cloud Translation API
Switch to English

XLA การโทรที่กำหนดเอง

เอกสารนี้อธิบายวิธีการเขียนและใช้ XLA "การโทรที่กำหนดเอง" การโทรแบบกำหนดเองช่วยให้คุณสามารถเรียกใช้โค้ดที่เขียนด้วยภาษาโปรแกรมเช่น C ++ หรือ CUDA จากโปรแกรม XLA

การโทรแบบกำหนดเองบน CPU

คุณสามารถสร้างคำสั่ง HLO ซึ่งแสดงถึงการโทรแบบกำหนดเองผ่านไคลเอนต์ API ของ XLA สิ่งนี้ไม่ได้เปิดเผยผ่าน TensorFlow ในขณะที่เขียน

ตัวอย่างเช่นรหัสต่อไปนี้ใช้การเรียกที่กำหนดเองเพื่อคำนวณ A[i] = B[i % 128] + C[i] บน CPU (แน่นอนคุณทำได้ - และควร! - ทำเช่นนี้กับ 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 นั้นค่อนข้างแตกต่างจากบน 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 = 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 CPU ของเรามีหน้าที่ในการจัดคิวงานบน GPU ที่นี่จะเปิดตัวเคอร์เนล CUDA แต่ก็สามารถทำอย่างอื่นได้เช่นเรียก cublas

buffers คืออาร์เรย์ของพอยน์เตอร์ที่อาศัยอยู่บนโฮสต์และแต่ละองค์ประกอบจะมีหน่วยความจำ point to device (เช่น GPU) พารามิเตอร์มาก่อนตามด้วยค่าเอาต์พุต นี้เป็นที่สะดุดตาแตกต่างจากการประชุม CPU โทรซึ่งมีสอง params, ins และ out เหตุผลหลักที่เราแยกความแตกต่างคือทำให้สามารถจัดการอินพุต / เอาต์พุตรูปทูเพิลได้อย่างมีประสิทธิภาพ ดูส่วนด้านล่าง

ดังตัวอย่างของ CPU เราได้กำหนดขนาดบัฟเฟอร์อินพุตและเอาต์พุตไว้ในการเรียกที่กำหนดเองของเรา อย่างไรก็ตามไม่เหมือนในเคสของ 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 มีการแสดงบัฟเฟอร์โพรโทคอลคุณจึงสามารถจัดเก็บโปรโตที่ต่อเนื่องกันนี้ไว้ภายใน opaque และยกเลิกการกำหนดค่าไว้ใน GPU ที่กำหนดเอง แต่ทราบว่าถึงแม้จะ xla::ShapeProto ไม่ได้เปลี่ยนบ่อยก็ไม่เปลี่ยนแปลง ตรวจสอบบันทึก git เพื่อดูว่าในอดีตมีการเปลี่ยนแปลงอย่างไร

การส่งผ่าน tuples ไปยังการโทรที่กำหนดเอง

พิจารณาการเรียกที่กำหนดเองต่อไปนี้

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 ในหน่วยความจำจะเหมือนกันใน CPU และ GPU แต่ก็มีการจัดการที่แตกต่างกันในรูปแบบการโทรที่กำหนดเองของ CPU และ GPU

เอาต์พุตทูเพิลเป็นบัฟเฟอร์ชั่วคราว

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

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

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

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

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

Tuples ในการเรียกที่กำหนดเองของ CPU

ในรหัส CPU เรามีฟังก์ชัน do_custom_call(const void** ins, void* out) ins เป็นอาร์เรย์ที่มีองค์ประกอบเดียวซึ่งชี้ไปที่ param0 บัฟเฟอร์ย่อยของ 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