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

Hướng dẫn này dành cho những kỹ sư hệ thống muốn XLA xuất những chương trình nhắm đến phần cứng của họ một cách hiệu quả. Hướng dẫn này không hướng dẫn 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 trình tăng tốc hoặc kiến trúc mới có thể triển khai để tạo một phần phụ trợ giúp chạy các chương trình học máy do XLA cung cấp. Quá trình 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 hoạt động hiện có từ một khung giao diện người dùng (chẳng hạn như PyTorch hoặc TensorFlow cho phần cứng mới).

Hầu hết các trường hợp triển khai sẽ rơi vào một trong các trường hợp sau:

  1. Kiến 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 hiện có.
  2. Phần cứng không giống CPU với phần phụ trợ LLVM hiện có.
  3. Phần cứng không giống CPU mà 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 cách sử dụ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ó phần phụ trợ LLVM cho phần cứng của họ, thì việc liên kết phần phụ trợ với LLVM tích hợp XLA là rất đơn giản. Ở chế độ JIT, phần phụ trợ CPU XLA sẽ phát ra mã cho CPU máy chủ. Để biên dịch trước, 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 chưa có phần phụ trợ LLVM 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 LLVM IR. Tuỳ thuộc vào tính chất của phần cứng, có thể sẽ phải thay đổi nhiều khía cạnh của quá trình tạo LLVM IR, nhưng bạn có thể chia sẻ nhiều mã với các phần phụ trợ hiện có.

Một 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, do đó, một số khía cạnh tạo mã của phần phụ trợ này chỉ dành riêng cho miền GPU. Các loại phần cứng khác, chẳng hạn như DSP như Hexagon (có phần phụ trợ LLVM ngược dòng), có thể dùng lại các phần của logic phát 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 mà không có phần phụ trợ LLVM

Nếu không thể sử dụng LLVM, thì tốt nhất là bạn nên triển khai một phần phụ trợ mới dành cho XLA đối với phần cứng mong muốn. Tuỳ chọn này đòi hỏi nhiều nỗ lự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 phương thức 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 được 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 tay điều khiển bộ nhớ thiết bị nhất định. Nói cách khác, mô-đun này giúp đóng gói hoạt động chuyển dữ liệu từ máy chủ lưu trữ sang thiết bị và ngược lại.