Phát triển phần phụ trợ mới cho XLA

Hướng dẫn này dành cho các kỹ sư hệ thống muốn XLA tạo ra các chương trình nhắm mục tiêu phần cứng một cách hiệu quả. Hướng dẫn này không chi tiết từng bước và giả định bạn có kiến thức về LLVM, Bazel và XLA.

XLA cung cấp một giao diện trừu tượng mà một cấu trúc hoặc trình tăng tốc mới có thể triển khai để tạo một phần phụ trợ nhằm chạy đầu ra của chương trình học máy do XLA cung cấp. Việc nhắm mục tiêu lại XLA sẽ đơn giản và dễ mở rộng hơn đáng kể so với việc triển khai mọi op hiện có từ khung giao diện người dùng (như PyTorch hoặc TensorFlow cho phần cứng mới).

Hầu hết các phương pháp triển khai sẽ thuộc một trong các trường hợp sau:

  1. Cấu trúc CPU hiện có chưa được XLA hỗ trợ chính thức, có hoặc không có phần phụ trợ LLVM.
  2. Phần cứng không giống CPU có phần phụ trợ LLVM hiện có.
  3. Phần cứng không giống CPU không có phần phụ trợ LLVM hiện có.

Trường hợp 1: Kiến trúc CPU hiện có chưa được XLA hỗ trợ chính thức

Trong trường hợp này, hãy bắt đầu bằng cách xem xét phần phụ trợ CPU XLA hiện có. XLA giúp bạn dễ dàng nhắm đến nhiều CPU bằng LLVM, vì điểm khác biệt chính giữa các phần phụ trợ XLA dành cho CPU là mã do LLVM tạo ra.

Nếu nhà cung cấp phần cứng có một phần phụ trợ LLVM cho phần cứng, thì bạn có thể dễ dàng liên kết phần phụ trợ đó với LLVM được xây dựng bằng XLA. Ở chế độ JIT, phần phụ trợ CPU XLA sẽ phát ra mã cho CPU máy chủ. Để biên dịch trước khi thực thi, xla::AotCompilationOptions có thể cung cấp bộ ba LLVM để định cấu hình cấu trúc mục tiêu.

Nếu không có phần phụ trợ LLVM hiện có nhưng có một loại trình tạo mã khác, bạn có thể sử dụng lại hầu hết phần phụ trợ CPU hiện có.

Trường hợp 2: Phần cứng không giống CPU với phần phụ trợ LLVM hiện có

Bạn có thể lập mô hình quy trình triển khai xla::Compiler mới trên các lớp xla::CPUCompiler và xla::GPUCompiler hiện có, vì các lớp này đã phát ra LLVM IR. Tuỳ thuộc vào bản chất của phần cứng, có thể sẽ phải thay đổi nhiều khía cạnh của thế hệ LLVM IR. Tuy nhiên, bạn có thể chia sẻ nhiều mã với các phần phụ trợ hiện có.

Ví dụ điển hình là phần phụ trợ GPU của XLA. Phần phụ trợ GPU nhắm mục tiêu đến một ISA không giống CPU, và do đó một số khía cạnh trong việc tạo mã chỉ dành riêng cho miền GPU. Các loại phần cứng khác, ví dụ: DSP như Hexagon (có phần phụ trợ LLVM ngược dòng), có thể sử dụng lại các phần của logic phát xạ LLVM IR, nhưng các phần khác sẽ là duy nhất.

Trường hợp 3: Phần cứng không giống CPU không có phần phụ trợ LLVM hiện có

Nếu không thể sử dụng LLVM, thì phương án tốt nhất là triển khai một phần phụ trợ mới cho XLA đối với phần cứng mong muốn. Tuỳ chọn này đòi hỏi nhiều công sức nhất. Sau đây là các lớp cần được triển khai:

  • StreamExecutor: Đối với nhiều thiết bị, không phải phương thức StreamExecutor nào cũng cần. Hãy xem các cách triển khai StreamExecutor hiện có để biết thông tin chi tiết.
  • xla::Compiler: Lớp này đóng gói hoạt động biên dịch phép tính HLO vào xla::Executable.
  • xla::Executable: Lớp này dùng để chạy một phép tính đã biên dịch trên nền tảng.
  • xla::TransferManager: Lớp này cho phép các phần phụ trợ cung cấp các cơ chế dành riêng cho nền tảng để tạo dữ liệu cố định XLA từ các ô điều khiển bộ nhớ thiết bị nhất định. Nói cách khác, phương thức này giúp đóng gói hoạt động truyền dữ liệu từ máy chủ đến thiết bị rồi ngược lại.