Xây dựng trình thông dịch Tail-call hiệu năng cao trong Rust

Trong tuần vừa qua, tôi đã thực hiện viết một trình thông dịch (interpreter) sử dụng kỹ thuật tail-call với từ khóa become, một tính năng mới vừa được thêm vào phiên bản nightly của Rust. Đây thực sự là một trải nghiệm thú vị và đáng ngạc nhiên, bởi máy ảo (VM) kết quả không chỉ hoạt động hiệu quả hơn phiên bản Rust trước đây mà còn vượt trội cả mã Assembly ARM64 do tôi tự viết tay.

Bài viết này là một bản tường thuật về việc triển khai một hệ thống tuy đơn giản nhưng không hề tầm thường, sử dụng các kỹ thuật dựa trên tail-call đang trở nên phổ biến gần đây. Đây cũng là bước tiếp theo trong việc khám phá việc giả lập hiệu năng cao CPU Uxn — một máy stack đơn giản phục vụ hệ sinh thái Hundred Rabbits.

Bối cảnh và vấn đề hiệu năng

Uxn là một máy stack có 256 lệnh instruction. Trong các trình giả lập đơn giản nhất, quy trình thực hiện thường là đọc một byte từ RAM tại bộ đếm chương trình (program counter), sau đó gọi hàm thực thi lệnh đó. Tuy nhiên, phương pháp này tồn tại nhược điểm lớn: các giá trị thường được lưu trong bộ nhớ thay vì thanh ghi (register), và việc lựa chọn nhánh để dispatch các lệnh tạo ra các chuỗi dự đoán khó khăn cho bộ vi xử lý.

Trước đó, để tối ưu hóa, tôi đã triển khai mã Assembly sử dụng kỹ thuật "threaded code" (mã luồng). Cách này lưu trữ toàn bộ trạng thái CPU trong các thanh ghi và kết thúc mỗi lệnh bằng một nhảy (jump) đến lệnh tiếp theo. Kết quả là tốc độ tăng trưởng đáng kể: nhanh hơn 40-50% trên ARM64 và gấp đôi trên x86-64. Tuy nhiên, cái giá phải trả là việc bảo trì khoảng 2000 dòng mã Assembly cực kỳ không an toàn (unsafe), từng dẫn đến lỗi ghi vượt vùng nhớ (out-of-bounds write) khó phát hiện.

Mục tiêu của chúng ta là đạt được hành vi tương tự như Assembly — trạng thái VM được lưu trong thanh ghi và dispatch ở cuối mỗi opcode — mà không cần viết từng lệnh một bằng Assembly.

Triển khai Tail-call trong Rust

Ý tưởng cốt lõi là sử dụng tail-call optimization (tối ưu hóa lời gọi đuôi). Chúng ta muốn trình biên dịch tạo ra lệnh nhánh br (branch to register) thay vì bl (branch-and-link), và quan trọng hơn là không cấp phát không gian persist nào trên stack. Nói cách khác, chúng ta cần một tail call thực thụ.

Với sự hỗ trợ của tính năng mới trong Rust, trình biên dịch đưa ra đảm bảo: khi thực hiện tail gọi một hàm, khung stack (stack frame) của người gọi sẽ được thay thế trực tiếp bởi khung stack của người được gọi, thay vì chồng lên nhau.

Về mặt triển khai, để tái sử dụng các opcode hiện có, chúng ta cần tái cấu trúc đối tượng Uxn ở đầu hàm, gọi hàm lệnh (ví dụ: inc), sau đó hủy cấu trúc nó khi gọi tác vụ tiếp theo. Mặc dù điều này tạo ra rất nhiều mã mẫu (boilerplate), nhưng chúng ta có thể loại bỏ nó bằng cách sử dụng Macro.

Macro này có thể trông khá phức tạp, nhưng nó cho phép khai báo các hàm mà vẫn giữ nguyên logic quan trọng. Đáng chú ý, toàn bộ quá trình này vẫn là 100% Safe Rust; thuộc tính #![forbid(unsafe_code)] vẫn được giữ nguyên. Trình biên dịch Rust làm rất tốt việc nội tuyến (inlining) và loại bỏ các hàm xuống các thao tác cốt lõi, loại bỏ hoàn toàn chi phí overhead của việc cấu trúc lại đối tượng UxnCore.

Đánh giá hiệu năng

Trên laptop sử dụng chip M1 (ARM64), tôi rất vui mừng khi trình thông dịch tail-call đã đánh bại mã Assembly viết tay của tôi trong cả hai bài kiểm tra benchmark. Điều này cho thấy trình biên dịch Rust đã tối ưu hóa mã bytecode cực kỳ tốt trên kiến trúc này.

Tuy nhiên, câu chuyện lại khác trên kiến trúc x86. Mặc dù phiên bản Rust tail-call hoạt động tốt hơn VM truyền thống, nó vẫn thua kém so với backend Assembly, đặc biệt trong bài kiểm tra microbenchmark Fibonacci.

Nguyên nhân nằm ở việc sinh mã (codegen). Trong khi các lệnh đơn giản như INC được biên dịch rất gọn gàng, các lệnh phức tạp hơn như ADD2 lại gặp vấn đề "register spilling" (tràn thanh ghi). Phiên bản Rust biên dịch ra mã dài hơn (121 byte so với 79 byte của Assembly) và phải đổ dữ liệu ra stack rồi khôi phục lại, làm giảm hiệu suất. Điều này có thể do tính năng tail-call trong Rust vẫn còn khá mới và chưa được tối ưu hóa hoàn toàn bởi backend LLVM.

Đối với WebAssembly (WASM), kết quả thậm chí còn thất vọng hơn. Trình thông dịch tail-call hoạt động chậm hơn nhiều so với bản build Rust native trên cả Firefox và Chrome. Có vẻ như các mẫu mã tạo ra Assembly tốt không ánh xạ tốt lên máy stack của WASM, và các trình JIT chưa đủ thông minh để tối ưu hóa chúng.

Kết luận

PR cho trình thông dịch tail-call đã được hợp nhất và triển khai trong bản phát hành 0.3.0 của dự án. Khi được bật, đây là tùy chọn mặc định trên hệ thống ARM64 và là lựa chọn thứ hai trên x86-64.

Đây là một bước tiến thú vị trong việc kết hợp sự an toàn của Rust với hiệu năng cao của kỹ thuật threaded code. Tôi rất mong chờ những cải tiến hơn nữa trong trình biên dịch để khắc phục các vấn đề về hiệu năng trên x86 và WebAssembly trong tương lai.