今年我將 PyPy RISC-V JIT Backend 送回上游的 PyPy 版本庫。今年 8 月發行的 PyPy v7.3.17 已經包含 RISC-V JIT Backend 的程式碼。這個 Side Project 我斷斷續續做了 8 年，總算在今年達成目標。我覺得值得寫一篇文章簡單地記錄我的心得。
簡介 PyPy 與 RPython
PyPy 是一個以 JIT Compiler（及時編譯器）加速執行 Python 程式碼的 Python 實作。根據它自己的 Benchmark，PyPy 大約比 CPython 3.7 快 4.4 倍。
PyPy 本身是以 RPython（Python 語言的子集）編寫的 Python 直譯器（Interpreter）。以 RPython 編寫的程式可以被「RPython Toolchain」編譯為機器碼直接執行。然而，即便把「PyPy 直譯器」編譯為機器碼，其執行 Python 程式的速度還是比「CPython 直譯器」慢。
為了讓 PyPy 變得比 CPython 快，PyPy 開發團隊實作了 Meta-Tracing JIT Compiler。為了方便解釋這個名詞，我把它拆解成三個部分：
JIT Compiler（及時編譯器）是指「執行環境（Run-time Environment）」在執行程式碼的時候，才將「程式碼片段」編譯為「機器碼」的編譯器。因為動態語言的型別資訊往往要等到執行期才能被確定下來，JIT Compiler 常見於各種動態語言的實作（例如：Javascript、Ruby 等等）。（備註：如果只使用 Ahead-of-Time Compiler，編譯器必須編譯很多特化版本再加上比較慢的通用版，但這又會增加機器碼大小，因此很難和直譯器或 JIT Compiler 競爭。）
JIT Compiler 根據編譯的「片段」分為兩個流派：Tracing JIT 與 Method JIT。Tracing JIT 是以直譯器蒐集反覆執行的「熱區」（可以是序列或迴圈），並將這些熱區編譯為機器碼。Method JIT 則是蒐集需要大量執行時間的函式，將被選中的函式編譯為機器碼。PyPy/RPython 使用的是 Tracing JIT。傳統觀點認為 Method JIT 有比較多程式碼可以被最佳化，因此其總體效能會比 Tracing JIT 好。但我個人持保留態度（這關乎 Benchmark 的特性與用來比較的 Method JIT 與 Tracing JIT），但那是另一個問題了。
Meta-Tracing JIT 是 Tracing JIT 的一個變型。如上所述，Tracing JIT 必須蒐集反覆執行的序列或迴圈。然而，典型的直譯器只有一個直譯器迴圈（Interpreter Loop）與若干個 Opcode Handler。如果直接以 Tracing JIT 蒐集直譯器的熱區，有可能所有的 Opcode Handler 都是熱區（但每個熱區都非常短）。如此一來，最好情況就只是相當於使用 C/C++ 語言編寫直譯器，並無法利用 JIT Compiler 特化型別的優勢。
Meta-Tracing JIT 的核心概念是讓 JIT Compiler 蒐集「被直譯的程式的熱區」而不是「直譯器的熱區」。為了達成這點，直譯器必須加上若干標記（Annotation）。透過這些標記，RPython 內建的 Meta-Tracing JIT 就能蒐集並編譯「被直譯的程式的熱區」。當然，PyPy 的 Python 直譯器已經標上這些標記。Meta-Tracing JIT 是 PyPy 與 RPython 的關鍵技術，不過其原理比較複雜，有興趣請參見 Tracing the meta-level: PyPy's tracing JIT compiler。
RISC-V JIT Backend
接下來描述我做的工作。JIT Backend 的工作就是把 Meta-Tracing 蒐集到的熱區編譯為機械碼。把一個 JIT Backend 移植到另一個平台基本是苦工。
RPython 的 JIT Backend 是一個簡單的 2-Pass Compiler：第一輪是 RPython 中間碼的暫存器分配器（Register Allocator）、第二輪把 RPython 中間碼轉譯為機器碼。RPython JIT Backend 程式碼的組織比較隨性。雖然有共用的 Utility Function，但是各個 Backend 的實作有很大的差異。因為我對 ARM/AArch64 比較熟，所以我是以 ARM/AArch64 Backend 為範本，並透過研究 ARM/AArch64 的程式碼與開發日誌理解每一個 Opcode 該做什麼。當 ARM/AArch64 Backend 有所不足，再以 X86 Backend 作為輔助。
在移植過程中，我體驗到測試導向開發（Test-Driven Development）的優點。整個 PyPy 與 RPython 是一個非常複雜的專案。它們有很多抽象層，所以很難理解每一個模組的功能。加上 RPython 只是一個由 PyPy 開發團隊定義的 Python 子集（更嚴格的說，它只是一個長得像是 Python 的語言），RPython Toolchain 的錯誤訊息有時候晦澀難解。因為這些門檻，開始做這個專案的前幾年，我一直在原地打轉。兩年前，我決定採用「測試導向開發」。先寫好 RISC-V 機器碼的組譯器（Assembler）、編寫 Test Case 驗收。接著依照 AArch64 Backend 的開發順序，一個一個移植 RPython 中間碼的編譯器。每移植完一個，就啟用一組 Test Case（RPython 有中間碼層級的 Test Case）。靠著「測試導向開發」我才能慢慢前進。我也在過程中學習到如何以 PyTest 與 Hypothesis 編寫單元測試。
值得一提的是：RPython 提供的 Test Case 有很高的涵蓋率。根據 RISC-V ABI，如果 Float Argument Register 用完了、但 Integer Argument Register 還沒用完，應該要把浮點數實際參數放在 Integer Argument Register。但是我一開始理解有誤，直接把浮點數實際參數放到 Stack。這也被 RPython 的 Test Case 抓到了（儘管它是一個 Platform-Independent Test）。
在 2023 年底，我利用年底假期一路推進到深水區。大約二週的時間裡，我一路加上 Function Call、Guard、Compiler Bridge、GC Write Barrier、GC Write Barrier (Array) 等等。
其中，GC Write Barrier (Array) 必須要在更新陣列元素的同時，標記 array[k * 128, ..., (k + 1) * 128 - 1] 是否被更新過，以便 Garbage Collector 在蒐集 Live Set 的時候可以一次跳過 128 個沒被更新的元素、減少 Garbage Collector 的暫停時間。這個 Opcode 把我帶回和學長討論 Garbage Collector 的時光。令人懷念。
到 2024 年初，我完成所有的 Opcode。終於可以用 RPython Toolchain 編譯整個 PyPy 直譯器。我寫了一封信到 PyPy 的 Mailing List：Contribute a RISC-V 64 JIT backend。依據 Matti Picus 的建議，我又跑了若干測試。除此之外，我也花了一段時間在 SiFive Unmatched 開發板上跑完 Benchmark。最後終於在今年 8 月送出 Pull Request 並於 8/14 被合併。有興趣的同好可以照著 Cross-Translating for RISC-V 建構帶有 RISC-V JIT Backend 的 PyPy。
後記
我很開心能在 2024 年完成這個斷斷續續執行 8 年多的 Side Project。這個 Side Project 讓我達成兩個心願：
- 第一個心願是親自從頭到尾將一個編譯器移植到新的平台。透過這個專案，我對 Compiler Backend 的細節（例如：Register Allocation、Function Call、Constant Pool、Trampoline 等等）有比較深刻的理解。當然我也藉由這個機會學習了 RISC-V 指令集與 RISC-V ABI。
- 第二個心願是貢獻一些 Patch 給 PyPy。自從我聽說 PyPy 之後，我一直在關注 PyPy 的 Blog。他們的 Blog 對我來說很有啟發性。PyPy 在很多技術岔路上都選擇了人跡罕至的道路。PyPy 在我的學生時代開拓了我的研究視野。這就是為什麼我一直無法放棄這個 Side Project。努力多年總算達成了！
最後，我想要特別感謝大學同學 X。我至今還記得 2009 年上半年的夜晚你在系館地下室熱情地向我介紹 PyPy 的情景。若你沒有向我介紹 PyPy，這一切都不會發生。謝謝你！
