pgrust: Tối ưu hóa JIT Compiler chỉ trong 5 micro giây nhờ AI

Trong lịch sử, biên dịch JIT (Just-In-Time) tốc độ cao từng được coi là một "nghệ thuật đen". Để viết một trình biên dịch JIT nhanh, bạn cần phải biết cách viết hợp ngữ (assembly). Điển hình là: hiện nay không có cơ sở dữ liệu sẵn sàng cho sản xuất (production-ready) nào tự sở hữu trình biên dịch JIT riêng. Tất cả chúng hoặc sử dụng LLVM, hoặc tạo ra mã C/C++. Cả hai lựa chọn này đều gặp vấn đề về thời gian biên dịch lâu, điều này hạn chế khả năng ứng dụng của chúng. Giờ đây, với sự hỗ trợ của AI, việc viết một trình biên dịch JIT với thời gian biên dịch nhanh bằng cách nhắm trực tiếp vào hợp ngữ trở nên dễ dàng hơn bao giờ hết. Đây cũng là một cơ hội cho các cơ sở dữ liệu mới cải tiến so với các hệ thống cũ. Khi xây dựng pgrust, ban đầu tôi nghĩ rằng việc triển khai trình biên dịch JIT sẽ rất khó khăn. Cuối cùng, tôi thấy nó dễ dàng hơn nhiều so với dự kiến nhờ sự hỗ trợ của AI, và đó cũng là một phần lý do tại sao pgrust lại nhanh đến vậy. Trình biên dịch JIT của pgrust biên dịch mã trong khoảng 5μs, cho phép chúng ta biên dịch JIT mọi truy vấn SQL, thay vì chỉ một tập con của chúng. Trong bài viết này, tôi sẽ hướng dẫn bạn cách tự xây dựng trình biên dịch JIT nhanh của riêng mình. Chúng ta sẽ xây dựng một công cụ biểu thức chính quy (regex) đơn giản sử dụng biên dịch JIT làm ví dụ.

Tại sao cần biên dịch JIT

Biên dịch JIT là phương pháp tạo ra mã đã biên dịch tại thời điểm chạy (runtime) hay còn gọi là "Just In Time". Khi được thực hiện đúng cách, nó có thể mang lại hiệu suất vượt trội, thường gấp 2-5 lần và đôi khi còn hơn thế nữa. Trường hợp sử dụng chính của biên dịch JIT là khi có những thông tin bạn chỉ thu thập được tại thời điểm chạy làm thay đổi đáng kể hành vi của chương trình. Điều này đặc biệt phổ biến với các trình thông dịch ngôn ngữ lập trình; chúng nhận mã để thực thi tại thời điểm chạy. Các trình biên dịch JIT cũng hữu ích trong các lĩnh vực ngoài ngôn ngữ lập trình, chẳng hạn như phân tích cú pháp dữ liệu. Đôi khi bạn không biết lược đồ (schema) của dữ liệu mình đang phân tích cho đến khi chạy, và JIT có thể giúp giải quyết vấn đề đó.

Để bắt đầu, hãy triển khai một công cụ biểu thức chính quy đơn giản. Để mọi thứ đơn giản, chúng ta chỉ hỗ trợ hai tính năng: chuỗi ký tự (literal strings) và lặp lại (ví dụ: regex *). Chúng ta cũng sẽ bỏ qua phần phân tích cú pháp (parser) và biểu diễn biểu thức chính quy dưới dạng các cấu trúc Rust đã được phân tích sẵn. Điều này có nghĩa là chúng ta sẽ có thể hỗ trợ các chuỗi như:

nhưng không hỗ trợ alternation (lựa chọn), lookbehind (nhìn ngược) hay bất kỳ tính năng tương tự nào.

Trong mã nguồn, điều này khá đơn giản. Chúng ta sẽ có 3 loại Node: node chuỗi ký tự, node lặp lại và node nối chuỗi (concatenation), là sự kết hợp của hai node. Kết quả sẽ trông như thế này:

Việc viết một trình thông dịch cho công cụ biểu thức chính quy của chúng ta cũng rất đơn giản:

Công cụ biểu thức chính quy này khá đơn giản. Nó chỉ dưới 20 dòng mã, nhưng hãy xem hiệu suất của nó ra sao. Để so sánh, chúng ta sẽ đối chiếu mã này với mã viết tay được tối ưu hóa riêng cho regex. Với ví dụ của chúng ta, chúng ta sẽ sử dụng regex b(an)*. Mã viết tay sẽ trông như sau:

(Có những cách để bạn tối ưu hóa mã này và làm cho nó nhanh hơn nhiều, nhưng với mục đích của chúng ta, nó đóng vai trò là một phép so sánh tốt)

Khi tôi chạy benchmark một vài ví dụ so sánh giữa hai phiên bản, tôi thấy phiên bản viết tay nhanh hơn 10-20 lần so với trình thông dịch. Rõ ràng là còn rất nhiều dư địa để cải thiện.

Bây giờ, hãy xem cách chúng ta có thể sử dụng biên dịch JIT để tạo ra một công cụ biểu thức chính quy tổng quát có hiệu suất ngang bằng với phiên bản viết tay.

Cách thực hiện biên dịch JIT

Có hai bước để biên dịch JIT mã nguồn. Đầu tiên, bạn tạo ra mã hợp ngữ cho đoạn mã bạn muốn chạy. Sau khi có mã, bạn đóng gói mã hợp ngữ đó vào một hàm mà bạn có thể gọi như bất kỳ đoạn mã nào khác trong chương trình của mình.

Để tạo mã hợp ngữ, chúng ta sẽ sử dụng một biến thể của phương pháp gọi là copy-and-patch. Ý tưởng là chúng ta có một loạt các mẫu (template) bằng hợp ngữ cho các thao tác khác nhau mà chúng ta muốn biên dịch JIT. Các mẫu này được gọi là "stencils". Khi muốn biên dịch JIT một thao tác, chúng ta lấy stencil tương ứng và thực hiện các điều chỉnh nhỏ dựa trên đặc thù của thao tác đó. Rất giống với việc điền vào một khuôn mẫu thực tế. Bằng cách nối các stencil đã điền này lại với nhau, chúng ta có thể xây dựng một chương trình tại thời điểm chạy có hiệu suất tương đương với phiên bản viết tay.

Đây là lộ trình chúng ta sẽ thực hiện: đầu tiên, chúng ta sẽ xem xét mã ARM64 mà chúng ta muốn tạo cho b(an)*. Sau đó, chúng ta sẽ biến các chuỗi lệnh lặp lại thành các stencil có thể tái sử dụng, viết một bộ phát (emitter) để điền và kết hợp các stencil đó từ AST của regex, và cuối cùng là sao chép các lệnh đã tạo vào bộ nhớ thực thi để Rust có thể gọi chúng như một hàm bình thường.

Để hướng dẫn bạn cách thức hoạt động, cách dễ nhất là bắt đầu với mã đã tạo và làm ngược lại về phía trình biên dịch JIT. Nhắc lại, chúng ta đang làm việc với regex “b(an)*”. Để thiết lập một số quyết định thiết kế:

Đối với trạng thái của chương trình, chúng ta sẽ sử dụng các thanh ghi sau:

Đối với các đầu vào của chương trình, chúng ta sẽ được truyền:

Mã ARM64 đã tạo

Bây giờ chúng ta đã xử lý xong phần đó, hãy cùng xem xét mã hợp ngữ đã tạo từng phần một. Đây là mã cụ thể trên macOS với ARM64. Đầu tiên, chúng ta có phần prologue (mở đầu), dùng để khởi tạo chương trình. Tất cả những gì nó làm là khởi tạo stack bằng cách thiết lập đỉnh và đáy của stack theo giá trị được truyền vào:

Tiếp theo, chúng ta có đoạn mã kiểm tra ký tự b. Nếu thấy một ký tự không phải là b, chúng ta nhảy đến một khối mã xử lý logic dự phòng (fallback). Nếu không, chúng ta tiến vị trí của mình trong chuỗi:

Tiếp theo, chúng ta có phần lặp lại (an)*. Đối với phần lặp lại, chúng ta cần thực hiện backtracking (quay lui). Nếu chúng ta quay lui tại đây, nghĩa là chúng ta nhảy ngay lập tức đến cuối vòng lặp. Điều đó có nghĩa là chúng ta cần lưu cả địa chỉ của lệnh sau vòng lặp và vị trí của chúng ta trong chuỗi vào stack.

Với thiết lập đó, giờ đây chúng ta có thể thực thi phần thân của vòng lặp. Nó sẽ kiểm tra các ký tự ‘a’ và ‘n’, và nếu tìm thấy, nó sẽ quay lại đầu vòng lặp, nhưng tại một vị trí chuỗi mới.

Bây giờ chúng ta đã vượt qua vòng lặp. Đây là nơi mà quá trình quay lui sẽ nhảy đến khi chúng ta thực hiện backtracking. Sau khi hoàn thành phần lặp lại, chúng ta đã ở cuối regex. Tất cả những gì cần làm bây giờ là kiểm tra xem chúng ta có đang ở cuối chuỗi hay không. Nếu đúng, chúng ta trả về 1 cho kết quả thành công. Nếu không, nghĩa là regex không khớp và chúng ta cần chạy logic thất bại để thực hiện dự phòng.

Và cuối cùng, chúng ta có logic dự phòng. Nó kiểm tra xem stack có trống hay không. Nếu trống, chúng ta trả về 0. Nếu không trống, chúng ta lấy cả địa chỉ dự phòng và vị trí chuỗi dự phòng ra khỏi stack, sau đó nhảy đến địa chỉ dự phòng đó.

Xây dựng các Stencil

Bây giờ bạn đã có cơ hội xem mã đã biên dịch, bạn sẽ bắt đầu hiểu cách trình biên dịch copy-and-patch hoạt động. Chúng ta có các tập hợp lệnh chung với chỉ những khác biệt nhỏ giữa chúng. Đối với mỗi khối hàm này, chúng ta có thể viết một hàm để tạo mã tương ứng. Mỗi hàm sẽ nhận các giá trị để sửa đổi mã. Ví dụ, một trong các đối số của stencil_char sẽ là ký tự trong regex cần so sánh. Chúng ta sẽ chèn trực tiếp ký tự đó vào mã máy.

Phần prologue rất đơn giản vì nó chỉ là một khối mã:

Đối với so sánh ký tự, chúng ta cần chèn ký tự đang so sánh và nơi cần nhảy đến cho logic dự phòng:

Đối với phần lặp lại, chúng ta có phần đầu của vòng lặp để đẩy vào stack và lệnh nhảy đến phần cuối:

Và sau đó chúng ta có các khối match (khớp) và fail (thất bại) khá gọn gàng:

Để đầy đủ, đây là các hàm hỗ trợ mà chúng ta đã sử dụng, giúp chèn dữ liệu cụ thể vào các lệnh:

Phát mã (Emitting Code)

Và đây là đoạn mã điều khiển nó:

Và đó là phần khó nhất! Cá nhân tôi thấy AI hữu ích nhất khi viết hợp ngữ. Kinh nghiệm chính của tôi với hợp ngữ là hoàn thành cuộc thi microcorruption CTF. Tôi chưa bao giờ thực sự tự viết hợp ngữ. Tôi sẽ rất chật vật để tìm ra chính xác các lệnh cần thiết và cách sửa đổi chúng để có được kết quả mong muốn. Với AI, tôi có thể cung cấp cho tác nhân lập trình của mình hình thái chung về cách tôi muốn trình biên dịch JIT hoạt động, và nó có thể xử lý rất nhiều chi tiết này thay cho tôi.

Nạp mã máy

Để hoàn thiện trình biên dịch, chúng ta cần thực sự nạp mã. Để làm điều này, chúng ta sẽ sử dụng mmap để cấp phát một khối bộ nhớ có thể đọc, ghi và thực thi. Sau đó, chúng ta sẽ sao chép mã vào bộ nhớ đó và chuyển đổi khối bộ nhớ đó thành một hàm mà chúng ta sẽ gọi:

Kết quả

Với tất cả những điều này đã hoàn tất, hãy so sánh hiệu suất của các triển khai khác nhau mà chúng ta đã xây dựng: