Lộ trình của StableHLO

Tại thời điểm viết bài, StableHLO đã sẵn sàng thay thế MHLO/HLO làm giao diện trình biên dịch. Mô hình này có thể được tạo bằng TensorFlow, JAX và PyTorch, có thể được sử dụng bởi XLA và một số trình bổ trợ PJRT của bên thứ ba, đồng thời có tất cả các tính năng công khai do MHLO/HLO cung cấp cũng như chức năng bổ sung.

Tài liệu này mô tả các bước tiếp theo cho dự án StableHLO, phân loại công việc đang diễn ra được phản ánh trong trình theo dõi lỗi và sắp xếp công việc này thành các sản phẩm dự kiến.

Các mốc quan trọng hiện tại

Các mốc quan trọng hiện tại của chúng tôi tuân theo hai xu hướng chính:

  1. Tối đa hoá lợi ích của StableHLO cho toàn bộ cộng đồng OpenXLA.
  2. Hợp nhất trải nghiệm nhà phát triển cho tất cả thành viên OpenXLA.
  • Ngừng sử dụng MHLO: Trong quý 4 năm 2024, chúng tôi đã bắt đầu tìm hiểu việc ngừng sử dụng MHLO nội bộ, di chuyển các lượt truyền hữu ích bao gồm cả việc chuẩn hoá và mẫu thư mục sang StableHLO. Sau khi quy trình di chuyển được chứng minh là không quan trọng trong nội bộ, chúng tôi dự định chia sẻ một RFC với tiến trình di chuyển bên ngoài sang StableHLO. Điều này có thể xảy ra vào Quý 1 năm 2025 và chúng tôi dự định sẽ dành nhiều thời gian và hỗ trợ các nhóm để di chuyển sang StableHLO trong Quý 1 năm 2025.
  • Di chuyển các hoạt động tối ưu hoá độc lập với phần cứng sang StableHLO: Theo xu hướng trên, chúng tôi muốn StableHLO là nơi tốt nhất để hợp nhất các hoạt động đơn giản hoá biểu đồ độc lập với phần cứng, để tất cả các trình bổ trợ PJRT, bao gồm cả các trình bổ trợ chuyển đổi từ StableHLO sang IR trình biên dịch không phải XLA, đều có thể thấy được lợi ích tối đa. Một phần của mục tiêu này liên quan đến việc hợp nhất các mẫu được sử dụng trong Google AI Edge, dự án JAX-Enzyme và các dự án khác đều nằm trong kho lưu trữ StableHLO. Một số hoạt động hợp nhất này đã bắt đầu, nhưng quy trình công việc sẽ chủ yếu được tiếp tục và hoàn tất trong quý 1 năm 2025.
  • Phân tách thành phần OpenXLA: Chúng tôi đã bắt đầu tạo các thành phần chuyên dụng trong openxla/xla cho HLO tương tự như cách thiết lập kho lưu trữ StableHLO (tham khảo), cũng như bắt đầu di chuyển tất cả phần phụ trợ OpenXLA đằng sau các trình bổ trợ PJRT. Ngoài ra, chúng tôi cũng đang đầu tư vào việc khắc phục các vấn đề nổi bật về trải nghiệm người dùng mà chúng tôi phát hiện thấy trong các trình bổ trợ PJRT này, bao gồm cả việc có thông tin liên lạc chính xác về phiên bản StableHLO trong các trình bổ trợ StableHLO, nhờ đó, các trình bổ trợ mới có thể sử dụng các tính năng mới ngay lập tức (tham khảo).
  • Đảm bảo thành phần kết hợp hoạt động từ đầu đến cuối: Trong quý 3 năm 2024, chúng tôi đã thêm thành phần kết hợp vào HLO, cho phép hỗ trợ ngăn xếp trình biên dịch đầy đủ cho các thành phần trừu tượng. Trong quý 4 năm 2024, chúng tôi đã hướng dẫn trình nội tuyến XLA về các thành phần kết hợp và thêm các lượt truyền trong HLO/StableHLO để nội tuyến các thành phần kết hợp không xác định với các thành phần phân ly của chúng. Chúng tôi hiện đang điều tra việc thêm các API JAX chuyên dụng để tạo thành phần kết hợp từ khung (các API PyTorch đã tồn tại), cũng như thêm tài liệu Colab về cách sử dụng thành phần kết hợp đúng cách, dự kiến hoàn tất vào quý 4 năm 2024.

Các cột mốc trước đây

Trong nửa đầu năm 2024, chúng tôi đã phát hành StableHLO v1.0. Phiên bản này đã triển khai các điểm cải tiến ưu tiên hàng đầu, bao gồm việc hoàn tất thông số kỹ thuật opset công khai, trình thông dịch tham chiếu, mở rộng khả năng tương thích ngược và xuôi để hỗ trợ triển khai trên thiết bị bằng cách kiểm thử đơn vị tương thích, khả năng mở rộng thông qua các thao tác tổng hợp, hỗ trợ tính năng linh động theo thông số kỹ thuật, phát hành một bộ dữ liệu kiểm thử đầy đủ với kết quả vàng được đánh giá, v.v.