Phát triển chương trình phụ trợ mới cho XLA

Hướng dẫn sơ bộ này dành cho những người dùng đầu tiên muốn dễ dàng nhắm mục tiêu lại TensorFlow vào phần cứng của họ một cách hiệu quả. Hướng dẫn không phải là bước-by-step và giả định kiến thức về LLVM , Bazel , và TensorFlow.

XLA cung cấp một giao diện trừu tượng mà một kiến ​​trúc hoặc bộ tăng tốc mới có thể triển khai để tạo phần phụ trợ để chạy đồ thị TensorFlow. Việc nhắm mục tiêu lại XLA sẽ đơn giản hơn và có thể mở rộng đáng kể so với việc triển khai mọi TensorFlow Op hiện có cho phần cứng mới.

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

  1. Hiện kiến trúc CPU vẫn chưa chính thức được hỗ trợ bởi XLA, có hoặc không có hiện LLVM backend.
  2. Phần cứng không giống CPU với phần mềm 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ó.

Tình huống 1: Kiến trúc CPU hiện tại chưa được XLA hỗ trợ chính thức

Trong tình huống này, hãy bắt đầu bằng cách nhìn vào hiện XLA CPU backend . XLA giúp dễ dàng nhắm mục tiêu lại TensorFlow đến các CPU khác nhau bằng cách sử dụng LLVM, vì sự 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. Google kiểm tra XLA cho kiến ​​trúc x64 và ARM64.

Nếu nhà cung cấp phần cứng có chương trình 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 được xây dựng bằng XLA rất đơn giản. Trong chế độ JIT, chương trình phụ trợ CPU XLA phát ra mã cho CPU chủ. Đối với phía trước-of-thời gian biên soạn, xla::AotCompilationOptions có thể cung cấp một LLVM triple để cấu hình các kiến trúc đích.

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

Tình huống 2: Phần cứng không giống CPU với phần mềm phụ trợ LLVM hiện có

Có thể mô hình hóa một mới xla::Compiler thực hiện trên hiện xla::CPUCompilerxla::GPUCompiler lớp học, vì đây đã phát ra LLVM IR. Tùy thuộc vào bản chất của phần cứng, có thể nhiều khía cạnh tạo LLVM IR sẽ phải được thay đổi, nhưng rất nhiều mã có thể được chia sẻ với các phần phụ trợ hiện có.

Một ví dụ tốt để sau là GPU backend của XLA. Phần phụ trợ GPU nhắm mục tiêu ISA không giống CPU và do đó một số khía cạnh của quá trình tạo mã của nó là duy nhất cho miền GPU. Các loại phần cứng khác, ví dụ như 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.

Tình huống 3: Phần cứng không giống CPU mà không có phần phụ trợ LLVM hiện có

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

  • StreamExecutor : Đối với nhiều thiết bị không phải tất cả các phương pháp của StreamExecutor là cần thiết. Xem hiện StreamExecutor triển khai để biết chi tiết.
  • xla::Compiler : Lớp này gói gọn việc biên soạn một tính toán HLO thành xla::Executable .
  • xla::Executable : Lớp này được sử dụng để khởi động một tính toán biên soạn trên nền tảng này.
  • xla::TransferManager : Lớp này cho phép backends để cung cấp cơ chế nền tảng cụ thể cho việc xây dựng XLA dữ liệu chữ từ xử lý bộ nhớ thiết bị nhất định. Nói cách khác, nó giúp đóng gói quá trình truyền dữ liệu từ máy chủ đến thiết bị và ngược lại.