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 cũng như 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 khung máy học (bao gồm PyTorch, TensorFlow và JAX) và nằm trong dự án OpenXLA – một hệ sinh thái gồm các công nghệ biên dịch nguồn mở dành cho máy học được các tổ chức phần mềm và phần cứng máy học hàng đầu hợp tác phát triển. 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ỏ chi phí trong thời gian chạy, kết hợp các hoạt động theo quy trình để giảm chi phí bộ nhớ và chuyên biệt các hình dạng tensor đã biết để cho phép truyền không đổi 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 sử dụng 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 tự động kết hợp để phù hợp với hiệu suất của các hoạt động tuỳ chỉnh ban đầu được kết hợp thủ công.
Cải thiện khả năng di chuyển. Giúp việc viết một phần phụ trợ mới cho phần cứng mới tương đối dễ dàng để một 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, vốn 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 thức hoạt động
Trình biên dịch XLA lấy các biểu đồ mô hình từ các khung học máy được xác định trong StableHLO và tổng hợp các biểu đồ đó thành các lệnh máy cho nhiều cấu trúc. StableHLO xác định một nhóm thao tác được tạo phiên bản (HLO = hoạt động cấp cao) cung cấp một lớp khả năng di chuyển giữa khung máy học và trình biên dịch.
Nói chung, quy trình biên dịch để chuyển đổi biểu đồ mô hình thành tệp thực thi được tối ưu hoá cho 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 không phụ thuộc vào mục tiêu, chẳng hạn như CSE, tổng hợp 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 quá trình 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 số liệu tính toán HLO đến một phần phụ trợ để tối ưu hoá ở cấp HLO tiếp theo, lần này là dựa trên thông tin và nhu cầu cụ thể về mục tiêu. Ví dụ: phần phụ trợ GPU có thể thực hiện các hợp nhất hoạt động mang lại lợi ích đặc biệt cho mô hình lập trình GPU và xác định cách phân vùng việc 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 tổ hợp các thao tác đó với 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ã theo 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 để tối ưu hoá, tạo mã và IR cấp thấp. Các phần phụ trợ này phát ra LLVM IR cần thiết để biểu thị hiệu quả tính toán HLO, 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 là mô-đun để 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 CPU.