JIT Compilation：理解與實現

本文主要介紹了基礎編譯技術中的 JIT Compilation 技術，以及如何使用 C++ 快速構建一個簡單的 JIT Compiler。

大約是在一年以前，“寫一篇文章介紹 JIT Compiler 是如何工作的” 這個想法就一直躺在我的 TODO List 中遲遲不能被點掉，而如今終於有時間將之付諸行動。那麼通過這篇文章，我希望能夠讓你瞭解到以下這些內容：

什麼是 JIT Compilation 技術？它有哪些特性？
如何使用 C++ 在不依賴任何框架的情況下實現一個 JIT Compiler？

而限於篇幅和主題範圍，本文將不會涉及以下這些內容：

如何編寫完備的 Interpreter / Compiler？
相關的高級編譯優化技術。

由於編寫 JIT Compiler 會涉及到從諸如 C/C++ 等高級編程語言、彙編、計算機體系結構，直到操作系統等多個方面的知識，因此這裏我將假設讀者已經具備這些領域相關的基礎知識，而當在文中實際涉及到相關內容時，我也會進行簡單的介紹。

在本文接下來將要闡述的例子中，考慮到完備性，以及爲了便於進行 Benchmark，我們會爲一個名爲 Brainfuck 的真實存在的編程語言實現一個簡單的 JIT Compiler。同時，我們也會爲其實現一個相應的 Interpreter，從而比較 JIT Compilation 與 Interpretation 這兩種方式在代碼整體執行效率上的差異。而關於 Interpreter 部分的具體實現細節，你可以參考例子所在倉庫中給出的源代碼，限於篇幅，本文將不做贅述。在我們正式開始之前，以下是你繼續閱讀所需要提前瞭解的一些事項：

代碼倉庫：https://github.com/Becavalier/brainfuck-jit-interpreter；
我們構建的 JIT Compiler 將以 X86-64 作爲目標平臺，其可以運行在 macOS 與 Linux 系統上；
由於 Compiler 實現部分的代碼較多，因此本文將有選擇性地進行介紹，完整代碼請參考上述倉庫。

好的，那讓我們開始吧。

Brainfuck 編程語言

Brainfuck 是一門從名字上來看就十分特殊的編程語言，它由 Urban Müller 於 1993 年創造。brainfuck 一詞本身是一個俚語詞彙，通常用來指帶那些超出人們理解的、非常複雜和罕見的事物，而 Brainfuck 這門語言便是如此。例如，以下這段 Brainfuck 代碼在正常執行後便會向控制檯輸出 “Hello, world!” 這幾個字符。

++++++++++[>+++++++>++++++++++>+++>+<<<<-]>++.>+.+++++++..+++.>++.<<+++++++++++++++.>.+++.------.--------.>+.>.

可以看到，通過肉眼識別代碼本身，我們根本無法得知整段程序的意圖，而這也正映射了 Brainfuck 語言其名稱的含義。雖然如此，但 Brainfuck 語言本身確是一門圖靈完備的極簡編程語言。Brainfuck 語言僅由 8 種不同的指令組成，所有由該語言編寫的程序均包含由這 8 種不同指令組成的不同指令序列。而程序在運行時，其包含的這些指令序列將被順序依次執行。除此之外，Brainfuck 的執行模型也十分簡單。除這 8 種不同的指令外，程序在執行時還會維護一個至少包含 30000 個單元的一維字節數組（後面我們將簡稱其爲 “紙帶”）。程序初始執行時，數組中的所有單元格均會被初始化爲數值 0，一個可以前後移動的 “數據指針” 將默認指向這個數組的第一個單元。而程序在運行時將會根據不同的指令來前後移動這個數據指針，並相應地更新或使用當前所指向單元格中的內容。關於上述提到的 8 種指令，它們所對應的字符和說明如下所示：

YUSazS

爲了加深理解，我們可以舉一個簡單的例子，比如下述這段 Brainfuck 代碼。

++[->+<]

這段代碼首先會將紙帶第一個單元格內的值連續遞增兩次（++），即變爲 2。隨後，[ 指令檢查到當前單元格內的值不爲 0（爲 2），因此繼續執行下一條指令。後續的四個指令 ->+< 會先將當前單元格內的值減一，接下來將數據指針向右移動到第二個單元格，然後將該單元格內的值加一，隨後再返回第一個單元格，如此往復循環。直達最後的 ] 指令判定第一個單元格內的值爲 0 時，程序結束運行。因此我們可以得知，這段程序主要用來更換紙帶前兩個單元格內的值。相應地，你也可以使用 Brainfuck Visualizer 來查看上述程序的完整動態執行過程。

JIT Compilation

在瞭解了我們的目標語言後，接下來讓我們一起看看 JIT Compilation 技術究竟是什麼？相信無論是做前端、後端，還是移動端，對於 “JIT” 一詞，你都肯定有所耳聞。JIT Compilation 的全稱爲 “Just-In-Time Compilation”，翻譯過來爲 “即時編譯”。其最顯著的特徵是代碼的編譯過程發生在程序的執行期間，而非執行之前。通常在編譯技術領域，我們會將 JIT 與 AOT 這兩種方式進行對比。AOT 編譯相信大家都十分熟悉，常見的比如：使用 Clang 對 C/C++ 代碼進行編譯、使用 Babel 編譯 ES6 代碼，甚至是將 JavaScript 代碼編譯爲專用於某一 JS 引擎的 IR（Intermediate Representation）等過程都可以被認作是 AOT 編譯的一種具體類型。而 JIT 與 AOT 之間的最大區別便是 “編譯過程發生的時間點”，對於 JIT 而言，其編譯過程發生在程序的運行時；而對 AOT 來說，編譯過程則發生在程序執行之前（通常爲構建時）。

傳統的 JIT 編譯器在實際動態生成機器碼前，會首先對原始代碼或其相應的 IR 中間代碼進行一系列的分析（profiling）。通過這些分析過程，編譯器能夠找到可以通過 JIT 編譯進行性能優化的 “關鍵代碼路徑”。而這裏的取捨重點在於：對這些代碼進行運行時優化而得到的性能提升收益，需要高於進行優化時所產生的性能開銷。在後面的文章中我們將會看到，對於我們的實例而言，這些開銷主要來自於代碼的運行時編譯，以及進行 OSR（On-Stack Replacement）的過程。而爲了便於理解，在本文後續的實例中，我們將不會實現傳統 JIT 所進行的代碼預分析過程。

另外需要注意的是，通常的 JIT 編譯器由於考慮到 “啓動延遲（startup time delay）” 的問題，因此一般會結合解釋器一起使用。JIT 編譯器所進行的代碼分析過程越精細、所實施的優化越多，其動態生成的機器代碼質量也會越高，但隨之而來的初始代碼執行延遲也會越大。而解釋器的加入便可使代碼的執行過程提前進行。而在此期間，JIT 編譯器也會同時對代碼進行分析和優化，並在特定的時刻再將程序的執行流程從解釋執行轉換到執行其動態生成的優化機器碼。因此，對於 JIT Compilation 這項技術而言，其實現方式需要取捨的一個重點是：在編譯時間和生成的代碼質量之間進行權衡。比如，JVM 便有着兩種可以選擇的 JIT 模式 —— client 與 server，其中前者會採用最小的編譯和優化選項以最大程度降低啓動延遲；而後者則會採用最大化的編譯和優化策略，同時犧牲程序的啓動時間。

實現細節

Ok, it’s time to showcase :)。首先聲明，我們爲 Brainfuck 語言實現的 JIT Compiler 只用於作爲本文內容的 POC，而並沒有考慮作爲生產版本的完備性，比如：exception-handling、thread-safe、profiling、assembly fine-tuning 等等。其次，接下來將要介紹的實現細節將着重聚焦於源代碼中的函數 bfJITCompile、函數 allocateExecMem，以及類 VM 這三個部分，這裏建議在繼續閱讀前，先自行大致瀏覽一下源代碼。

就如同上面我們所說的那樣，JIT Compilation 的代碼編譯過程發生在程序的運行時，因此從源代碼中也可以看到，我們通過用戶在運行解釋器程序時所提供的不同參數（–jit）來決定是採用 JIT 編譯執行，還是直接解釋執行。而對於 “JIT 編譯執行” 這種方式來說，其流程可大致總結爲：

讀入源代碼（包含 ASCII 形式的指令序列）；
調用 bfJITCompile 函數，將源代碼編譯爲機器碼；
調用 allocateExecMem 函數，將機器碼動態分配在可執行的內存段上；
調用 VM::exec 函數，通過 OSR 轉移執行流程；
代碼執行完畢後再次轉移回主流程；
執行一些清理善後工作。

接下來，我們將重點介紹上述流程中的第二、三及第四項的具體實現細節。

編譯機器碼

在這一步中，我們會將程序啓動時輸入的 Brainfuck 源代碼中的所有指令字符全部 “提取” 出來，並直接按順序爲其生成相應的機器碼版本的二進制代碼。這些生成的二進制機器碼集合將被存放在一個 std::vector 對象中以備後續使用。爲了簡化機器碼的生成過程，我們簡單地通過 switch 語句識別出指令對應的字符，並 “返回” 該指令對應的 X86-64 二進制機器碼。而這些返回的機器碼也將被直接 “拼接” 到用於存放機器碼集合的 Vector 容器中。

這裏需要注意的是，對於這些返回的二進制機器碼，由於其中可能含有引用的相對地址信息（RIP-Relative），因此在被實際存放到 Vector 容器之前，我們還需要通過諸如 _relocateAddrOfPrintFunc 等方法來對這些二進制機器碼進行 “地址重定位” 處理。通過這些方法，我們能夠準確計算出這些相對地址的實際信息，並對它們進行修正。

首先，在 bfJITCompile 函數的定義中我們可以找到如下這段代碼。通過這段代碼，我們將 Brainfuck 執行模型中的 “數據指針” 其地址存放在了寄存器 rbx 中，這樣後續我們便可以通過修改或使用該寄存器中的值來控制數據指針的位置，或者讀取、修改當前數據指針所指向紙帶單元格中的內容。這裏代碼中的 “/* mem slot */” 表示該註釋所在位置的內容將在編譯時被替換爲實際引用的內存地址。而這個地址將來自於 bfState::ptr 的值在經過函數 __resolvePtrAddr_ 處理後返回的小端（_little-endian_）格式地址。

// prologue.std::vector<uint8_t> machineCode {
  // save dynamic pointer in %rbx.
  0x48, 0xbb, /* mem slot */};// ...

接下來，隨着不斷讀入的指令字符，bfJITCompile 函數便可以依次將這些指令轉換爲其對應的機器碼版本。對於 “+ -> <” 這四個指令來說，它們所對應的機器指令只需要通過操作我們先前存放在 rbx 寄存器中的數據指針的地址值，便可完成對 Brainfuck 抽象機器的狀態改變。比如以 “+” 指令爲例，我們可以找到如下這段代碼：

// ...case '+': {
  for (n = 0; *tok == '+'; ++n, ++tok);
  const auto ptrBytes = _resolvePtrAddr(ptrAddr);
  std::vector<uint8_t> byteCode { 
    0x80, 0x3, static_cast<uint8_t>(n),  // addb $0x1, (%rbx)
  };
  _appendBytecode(byteCode, machineCode);
  --tok;
  break;} // ...

在這段代碼中我們首先使用了一個很容易想到的優化策略，那就是當遇到連續的 “+” 指令時，相較於爲每一個出現的 “+” 指令都生成相同的、重複的機器碼，我們可以選擇首先計算遇到的連續出現的 “+” 指令的個數，然後再通過一條單獨的彙編指令 addb $N, (%rbx) 來將這多個 “+” 指令所產生的狀態變更一次性完成。相同的方式還可以被應用到其餘的三種指令，它們分別對應數據指針所指向單元格內值的改變，以及數據指針本身的值的改變。

而對於 “,” 及 “.” 指令來說，由於它們涉及 IO 操作，因此這裏對應的機器碼將涉及對操作系統調用（System Call）的調用過程。操作系統調用需要遵循特定的調用慣例（Calling Convention），比如通常來說，寄存器 rax 用於存放系統調用號、rdi 用於存放第一個參數、rsi 用於存放第二個參數，以及 rdx 用於存放第三個參數等等。同時，macOS 與 Linux 操作系統下的系統調用號也並不相同，這裏我們通過預處理指令來進行區分。

// ...case ',': {
  /**
    movl $0x2000003, %eax
    movl $0x0, %edi
    movq %rbx, %rsi
    movl $0x1, %edx
    syscall
  */
  std::vector<uint8_t> byteCode { #if __APPLE__
    0xb8, 0x3, 0x0, 0x0, 0x2,#elif __linux__
    0xb8, 0x0, 0x0, 0x0, 0x0,#endif
    0xbf, 0x0, 0x0, 0x0, 0x0,
    0x48, 0x89, 0xde,
    0xba, 0x1, 0x0, 0x0, 0x0,
    0xf, 0x5,
  };
  _appendBytecode(byteCode, machineCode);
  break;}// ...

最後，對於 “[” 和 “]” 指令，其實現邏輯會稍微有些複雜。以 “[” 指令爲例，如下代碼所示。在這裏，將 “[” 指令的語義邏輯直接映射到彙編代碼是十分簡單的，其邏輯是：判斷當前數據指針所指向單元格的值是否爲 0。若爲 0，則執行流程跳轉到後續與其配對的 “]” 指令的後一個指令；否則繼續執行下一條指令。因此，我們這裏直接使用 cmpb 彙編指令來判斷以寄存器 rbx 中的值作爲地址時，其對應內存位置的值是否爲 0。若爲 0，則使用 je 彙編指令跳轉到後續的指令位置，否則直接執行下一條指令。代碼中對 “後續指令地址” 的使用將會在與其配對的 “]” 指令處理流程中對其進行重定向處理。因此，這裏我們將使用連續四個字節的 0x0 值進行佔位。另外需要知道的是，爲了簡化實現，這裏我們將固定使用 “near jump” 模式。

// ...case '[': {
  /*
    cmpb $0x0, (%rbx)
    je <>
  */
  std::vector<uint8_t> byteCode { 
    0x80, 0x3b, 0x0,
    0xf, 0x84, 0x0, 0x0, 0x0, 0x0, /* near jmp */
  };
  // record the jump relocation pos.
  _appendBytecode(byteCode, machineCode);
  jmpLocIndex.push_back(machineCode.size());
  break;}// ...

至此，我們便完成了機器指令的動態編譯工作。通過這個階段，我們的程序可以將輸入的 Brainfuck 指令字符序列轉換成對應的平臺相關的二進制機器碼。你可以在 bfJITCompile 函數的最後看到如下這樣一段用來收尾的代碼。這段代碼主要用於在程序退出前輸出 stdout 緩存區中的內容，並重置 rip 寄存器的值，以將程序執行流程退回到 C++ 代碼中。後續我們還將回顧這部分內容。

// epilogue. // mainly restoring the previous pc, flushing the stdout buffer./**
  cmpq $0, %r11
  je 8
  callq <print>
  jmpq *(%rsp)
  */std::vector<uint8_t> byteCode {
  0x49, 0x83, 0xfb, 0x0,
  0x74, 0x8,
  0xe8, /* mem slot */
  0xff, 0x24, 0x24,};

可執行內存分配

接下來，我們將關注點從 “如何動態生成機器碼” 移到 “如何動態執行機器碼” 這個問題上。關於這部分實現可以參考名爲 allocateExecMem 的函數，相關如下代碼所示。

// ...uint8_t* allocateExecMem(size_t size) {
  return static_cast<uint8_t*>(
    mmap(
      NULL,
      size, 
      PROT_READ | PROT_WRITE | PROT_EXEC, 
      MAP_PRIVATE | MAP_ANONYMOUS, 
      -1,
      0));}// ...

在這個函數的定義中，我們調用名爲了 mmap 的函數，而這個函數便是支持 “動態執行機器碼” 的關鍵所在。mmap 函數是一個由 C 標準庫提供的系統調用，通過該函數，我們可以在當前進程的 VAS（Virtual Address Space）中創建一個映射。這個映射可以指向一個具體的文件、或者是一個匿名空間。關於 mmap 函數，我們最爲熟知的一種使用方式便是在爲目標文件分配虛擬頁時，操作系統會使用該函數將頁表條目指向目標文件中的適當位置。而在這裏，我們則需要利用該函數創建不指向任何實際文件的 “匿名空間”，並將我們在上一步中編譯得到的二進制機器碼連續地放入到這段內存空間中。

不僅如此，通過爲 mmap 函數的第三個參數指定 PROT_EXEC 屬性，我們可以將這段申請的匿名內存空間標記爲 “可執行”。這意味着，存放在這段內存空間中的機器指令可以被 CPU 執行。而關於該函數其他參數的詳細配置信息，你可以參考這裏找到更多答案。allocateExecMem 函數的實際調用過程則被放置在了 VM 類的構造函數中，在這裏我們通過 RAII 將資源的分配與銷燬進行了簡單的封裝。

OSR（On-Stack Replacement）

當編譯生成的二進制機器碼被放入可執行的匿名內存段後，接下來的重點便是：_如何將程序的指令執行流程轉移至這段內存的起始位置處？_關於這部分實現，我們需要藉助 Clang/GCC 編譯器提供的 “C++ 內聯彙編” 功能。你可以在 VM::exec 函數的實現中找到答案。這段代碼如下所示：

// ...void exec() {
  // save the current %rip on stack (by PC-relative).
  // %r10 - stdout buffer entry.
  // %r11 - stdout buffer counter.
  asm(R"(
    pushq %%rax
    pushq %%rbx
    pushq %%r10
    pushq %%r11
    pushq %%r12
    movq %1, %%r10
    xorq %%r11, %%r11
    lea 0x9(%%rip), %%rax 
    pushq %%rax
    movq %0, %%rax
    addq %2, %%rax
    jmpq *%%rax 
  )":: "S" (mem), "m" (stdoutBuf), "D" (prependStaticSize));

  // clean the stack.
  asm(R"(
    addq $8, %rsp
    popq %r12
    popq %r11
    popq %r10
    popq %rbx
    popq %rax  )");}// ...

在這段代碼中，我們使用了兩次 asm 彙編指令。其中，第一次內聯彙編的目的主要是爲了將程序的執行流程轉移到我們之前動態編譯生成的機器碼上。這裏前 5 行對 push 指令的調用過程主要用於將這些寄存器中的值存放到棧上，以保護此刻的寄存器狀態。而這些值將會在程序的執行流程返回到 C++ 代碼後再被重新恢復。第 6 行的 movq 指令將 stdout buffer 的首地址存放到了寄存器 r10 中，這個 buffer 將用於緩存通過 “.” 指令輸出的字符內容，以減少系統調用的實際調用次數，提升性能。接下來的第 8-9 行，我們將正常 C++ 代碼執行流程的下一條指令其地址存放到了棧上，以便後續能夠從動態執行流程中正常返回。第 10-11 行，我們正確地設置了匿名可執行內存段的地址以及相應的偏移位置（跨過了靜態的 subroutine 定義部分）。最後一行，通過 jmpq 指令，我們讓 CPU 的執行流程跳轉到以 rax 寄存器中的值作爲內存地址的位置，即包含我們將要執行的第一條動態指令的位置。

至此，從 C++ 代碼到動態指令的執行轉移流程便完成了。而當動態生成的指令全部執行完畢後，我們需要通過類似的方式再將執行流程轉移回正常的 C++ 代碼中。還記得我們在 “編譯機器碼” 這一小節最後提到的那小段 “epilogue” 彙編代碼嗎？如果返回去查看，你會發現在這段代碼的最後一條指令中，我們使用了 jmpq *(%rsp) 指令，這條指令將會把 CPU 的執行流程轉移到以當前進程棧底存放的那個 qword 值作爲地址的內存位置上。而這個值，便是我們在上一步中存放的 C++ 代碼的返回地址。當執行流程返回到 C++ 代碼後，我們遇到了第二個 asm 彙編指令。通過這段指令，我們可以清理棧上的內容並同時恢復相關寄存器的狀態。到這裏，程序的執行流程便基本結束了。

讓我們將目光再移回到本小節的主題 “OSR” 上來。OSR 的全稱爲 “On-Stack Replacement”。藉助 Google，我們可以找到對它的一個定義，如下所示：

On-stack-replacement (OSR) describes the ability to replace currently executing code with a different version, either a more optimized one (tiered execution) or a more general one (deoptimization to undo speculative optimization).

實際上，對於 OSR 我們可以將其簡單理解爲 “從一個執行環境到另一個執行環境的轉換過程”。比如在我們的實例中，VM::exec 函數在執行時，它會將執行環境從 C++ 代碼轉移至動態生成的機器碼，最後再以同樣的方式轉移回來。而這樣的執行環境轉換便可被視爲 OSR 的過程。下圖是對上述 OSR 過程的一個形象展示。

Benchmark

至此，我們已經介紹完了 Brainfuck JIT Compiler 幾個關鍵點的實現細節。那現在讓我們來看看這個粗糙版的 JIT 編譯器其性能如何？項目的源代碼中提供了兩組測試，分別用於測試 “IO 密集型” 和 “計算密集型” 這兩個場景。一組測試結果如下所示：

IO 密集型 case：

Benchmark for 10 seconds: (higher score is better)
   12950 interpreter
   35928 jit (win)

計算密集型 case：

Benchmark Result: (lower time is better)
    13.018s interpreter
     0.885s jit (win)

可以看到，總體結果還算不錯。對於 IO 密集型的測試用例，JIT Compilation 相比單純的 Interpretation 可以帶來將近 3 倍的性能提升。而對於計算密集型場景來說，JIT 帶來的性能提升便十分可觀了。在我們的 “打印曼德布洛特集合” 的測試用例中，使用 JIT Compilation 相較於 Interpretation 可以帶來將近 15 倍的性能提升。當然，鑑於我們並沒有採用更加完備的測試集合及測試方案，這些測試用例結果僅供參考。

JIT Compilation：理解與實現

Brainfuck 編程語言

JIT Compilation

實現細節

編譯機器碼

可執行內存分配

OSR（On-Stack Replacement）

Benchmark

更多信息

Interpretation 之殤

JIT 實現策略與方式

Brainfuck 編程語言

JIT Compilation

實現細節

編譯機器碼

可執行內存分配

OSR（On-Stack Replacement）

Benchmark

更多信息

Interpretation 之殤

JIT 實現策略與方式

猜你喜歡