Giúp bảo vệ Great Barrier Reef với TensorFlow trên Kaggle Tham Challenge

Kiến trúc XLA

Tại sao chúng tôi xây dựng XLA?

Chúng tôi đã có một số mục tiêu để XLA làm việc với TensorFlow:

  • Cải thiện tốc độ thực thi. Biên dịch các đồ thị con để giảm thời gian thực hiện của các Ops tồn tại ngắn hạn để loại bỏ chi phí khỏi thời gian chạy TensorFlow, kết hợp các hoạt động pipelined để giảm chi phí bộ nhớ và chuyên biệt hóa các hình dạng tensor đã biết để cho phép truyền liên tục tích cực hơn.

  • Cải thiện việc sử dụng bộ nhớ. Phân tích và lập lịch sử dụng bộ nhớ, về nguyên tắc loại bỏ nhiều bộ đệm lưu trữ trung gian.

  • Giảm sự phụ thuộc vào các Ops tùy chỉnh. Loại bỏ nhu cầu về nhiều Hoạt động tùy chỉnh bằng cách cải thiện hiệu suất của các Hoạt động cấp thấp được hợp nhất tự động để phù hợp với hiệu suất của các Hoạt động tùy chỉnh đã được hợp nhất bằng tay.

  • Giảm dấu chân di động. Loại bỏ thời gian chạy TensorFlow bằng cách biên dịch trước thời hạn biểu đồ con và tạo ra một cặp tệp đối tượng / tiêu đề có thể được liên kết trực tiếp với một ứng dụng khác. Kết quả có thể làm giảm dấu chân cho suy luận di động theo một số bậc lớn.

  • Cải thiện tính di động. Làm cho nó tương đối dễ dàng để viết một chương trình phụ trợ mới cho phần cứng mới, tại thời điểm đó, một phần lớn các chương trình TensorFlow sẽ chạy không được sửa đổi trên phần cứng đó. Điều này trái ngược với cách tiếp cận chuyên biệt hóa các Ops nguyên khối riêng lẻ cho phần cứng mới, đòi hỏi các chương trình TensorFlow phải được viết lại để tận dụng các Ops đó.

XLA hoạt động như thế nào?

Ngôn ngữ đầu vào cho XLA được gọi là "HLO IR", hay chỉ là HLO (Hoạt động mức cao). Ngữ nghĩa của HLO được mô tả trên Operation Semantics trang. Đó là thuận lợi nhất để nghĩ về HLO như một IR trình biên dịch .

XLA nhận các đồ thị ("tính toán") được xác định trong HLO và biên dịch chúng thành các lệnh máy cho các kiến ​​trúc khác nhau. XLA là mô-đun theo nghĩa là nó rất dễ dàng để khe trong một backend thay thế cho nhắm mục tiêu một số kiến trúc HW cuốn tiểu thuyết . Phần phụ trợ CPU cho x64 và ARM64 cũng như phần phụ trợ GPU NVIDIA nằm trong cây nguồn TensorFlow.

Sơ đồ sau đây cho thấy quá trình biên dịch trong XLA:

XLA đi kèm với một vài tối ưu hóa và đèo phân tích đó là mục tiêu độc lập, chẳng hạn như CSE , mục tiêu độc lập hoạt động hợp nhất, và đệm phân tích cho việc cấp phát bộ nhớ thời gian chạy cho việc tính toán.

Sau bước không phụ thuộc vào đích, XLA sẽ gửi tính toán HLO đến một chương trình phụ trợ. Phần phụ trợ có thể thực hiện tối ưu hóa cấp HLO hơn nữa, lần này có lưu ý đến thông tin và nhu cầu cụ thể của mục tiêu. Ví dụ: phụ trợ GPU XLA có thể thực hiện kết hợp hoạt động có lợi cụ thể cho mô hình lập trình GPU và xác định cách phân vùng tính toán thành các luồng. Ở giai đoạn này, các chương trình phụ trợ cũng có thể khớp với một số hoạt động hoặc kết hợp nhất định của chúng với các lệnh gọi thư viện được tối ưu hóa.

Bước tiếp theo là tạo mã mục tiêu cụ thể. CPU và GPU backends kèm XLA sử dụng LLVM cho cấp thấp IR, tối ưu hóa, và mã thế hệ. Các phần phụ trợ này phát ra LLVM IR cần thiết để biểu diễn tính toán XLA HLO một cách hiệu quả, sau đó gọi LLVM để phát ra mã gốc từ LLVM IR này.

Phần phụ trợ GPU hiện hỗ trợ GPU NVIDIA thông qua phần phụ trợ LLVM NVPTX; chương trình phụ trợ CPU hỗ trợ nhiều ISA CPU.