Trang này được dịch bởi Cloud Translation API.
Switch to English

Phát triển phần 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à từng bước 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 trình tăng tốc mới có thể thực hiện để tạo một phụ trợ để chạy các biểu đồ TensorFlow. Nhắm mục tiêu lại XLA nên đơn giản và có khả năng mở rộng đáng kể hơn 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 tình huống sau:

  1. Kiến trúc CPU hiện tại chưa được XLA hỗ trợ chính thức, có hoặc không có phụ trợ LLVM hiện có.
  2. Phần cứng không giống CPU với phụ trợ LLVM hiện có.
  3. Phần cứng không giống CPU mà không có phụ trợ LLVM hiện có.

Kịch bản 1: Kiến trúc CPU hiện tại chưa được XLA hỗ trợ chính thức

Trong kịch bản này, bắt đầu bằng cách xem xét phụ trợ CPU XLA hiện có. XLA giúp dễ dàng nhắm mục tiêu lại TensorFlow cho 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ụ trợ XLA cho CPU là mã do LLVM tạo ra. Google kiểm tra XLA cho các kiến ​​trúc x64 và ARM64.

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ật đơn giản để liên kết phần phụ trợ với LLVM được xây dựng với XLA. Trong chế độ JIT, phụ trợ CPU XLA phát mã cho CPU chủ. Để biên dịch trước thời hạn, xla::AotCompilationOptions có thể cung cấp bộ ba LLVM để định cấu hình kiến ​​trúc đích.

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

Kịch bản 2: Phần cứng không giống CPU với phụ trợ LLVM hiện có

Có thể mô hình hóa một triển khai xla::Compiler trên các xla::CPUCompilerxla::GPUCompiler , vì các lớp nà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 thay đổi, nhưng rất nhiều mã có thể được chia sẻ với các phụ trợ hiện có.

Một ví dụ điển hình để làm theo là phần phụ trợ GPU của XLA. Phần cuối của GPU nhắm vào một ISA không giống CPU, và do đó một số khía cạnh trong việc 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ụ DSP như Hexagon (có 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.

Kịch bản 3: Phần cứng không giống CPU mà không có 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 một 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 thực hiện như sau:

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