XLA (Đại số tuyến tính tăng tốc) là một trình biên dịch học máy (ML) giúp tối ưu hoá đại số tuyến tính, giúp cải thiện tốc độ thực thi và mức sử dụng bộ nhớ. Trang này cung cấp thông tin tổng quan ngắn gọn về mục tiêu và cấu trúc của trình biên dịch XLA.
Mục tiêu
Hiện nay, XLA hỗ trợ một số giao diện người dùng của khung ML (bao gồm PyTorch, TensorFlow và JAX) và là một phần của dự án OpenXLA – một hệ sinh thái gồm các công nghệ biên dịch nguồn mở cho ML được các tổ chức phần cứng và phần mềm hàng đầu phát triển cùng nhau. Trước khi tạo dự án OpenXLA, XLA đã được phát triển bên trong dự án TensorFlow, nhưng các mục tiêu cơ bản vẫn không thay đổi:
Cải thiện tốc độ thực thi. Biên dịch các đồ thị con để giảm thời gian thực thi của các hoạt động ngắn hạn và loại bỏ hao tổn trong thời gian chạy, hợp nhất các hoạt động quy trình để giảm mức hao tổn bộ nhớ và chuyên biệt hoá các hình dạng tensor đã biết để cho phép truyền hằng số linh hoạt hơn.
Cải thiện mức sử dụng bộ nhớ. Phân tích và lên lịch sử dụng bộ nhớ, loại bỏ nhiều vùng đệm lưu trữ trung gian.
Giảm sự phụ thuộc vào hoạt động tuỳ chỉnh. Loại bỏ nhu cầu phải có nhiều hoạt động tuỳ 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 kết hợp tự động để phù hợp với hiệu suất của các hoạt động tuỳ chỉnh ban đầu được hợp nhất theo cách thủ công.
Cải thiện khả năng di chuyển. Hãy tạo điều kiện để tương đối dễ dàng viết một phần phụ trợ mới cho phần cứng mới, để phần lớn các mô hình học máy có thể chạy mà không bị sửa đổi trên phần cứng đó. Điều này trái ngược với phương pháp chuyên về các hoạt động nguyên khối riêng lẻ cho phần cứng mới, trong đó yêu cầu phải viết lại các mô hình để sử dụng các hoạt động đó.
Cách hoạt động
Trình biên dịch XLA lấy biểu đồ mô hình từ các khung học máy được xác định trong StableHLO và biên dịch các biểu đồ đó thành hướng dẫn máy cho nhiều cấu trúc. StableHLO xác định một tập hợp thao tác được tạo phiên bản (HLO = các thao tác cấp cao) cung cấp một lớp khả năng di chuyển giữa khung ML và trình biên dịch.
Nhìn chung, quá trình biên dịch chuyển đổi biểu đồ mô hình thành tệp thực thi tối ưu hoá mục tiêu bao gồm các bước sau:
XLA thực hiện một số lần truyền phân tích và tối ưu hoá tích hợp sẵn trên biểu đồ StableHLO độc lập với mục tiêu, chẳng hạn như CSE, hợp nhất hoạt động độc lập mục tiêu và phân tích vùng đệm để phân bổ bộ nhớ thời gian chạy cho việc tính toán. Trong giai đoạn tối ưu hoá này, XLA cũng chuyển đổi phương ngữ StableHLO thành một phương ngữ HLO nội bộ.
XLA gửi phép tính HLO đến một phần phụ trợ để tối ưu hoá thêm ở cấp độ HLO, lần này là tập trung vào thông tin và nhu cầu cụ thể theo mục tiêu. Ví dụ: phần phụ trợ GPU có thể thực hiện việc hợp nhất hoạt động mang lại lợi ích cụ thể cho mô hình lập trình GPU và xác định cách phân chia hoạt động tính toán vào các luồng. Ở giai đoạn này, các phần phụ trợ cũng có thể so khớp mẫu một số thao tác hoặc kết hợp các thao tác đó với các lệnh gọi thư viện được tối ưu hoá.
Sau đó, phần phụ trợ sẽ thực hiện việc tạo mã cho từng mục tiêu cụ thể. Các phần phụ trợ CPU và GPU đi kèm với XLA sử dụng LLVM để IR, tối ưu hoá và tạo mã cấp thấp. Các phần phụ trợ này phát ra LLVM IR cần thiết để biểu thị việc tính toán HLO một cách hiệu quả, sau đó gọi LLVM để phát mã gốc từ LLVM IR này.
Trong quy trình này, trình biên dịch XLA được mô-đun hoá có thể dễ dàng đưa vào một phần phụ trợ thay thế để nhắm đến một số cấu trúc HW mới. Phần phụ trợ GPU hiện hỗ trợ GPU NVIDIA thông qua phần phụ trợ LLVM NVPTX. Phần phụ trợ CPU hỗ trợ nhiều ISA của CPU.