JIT Code Generation代碼生成
一．表達式編譯
代碼生成（Code Generation）技術廣泛應用于現代的數據系統中。代碼生成是將用戶輸入的表達式、查詢、存儲過程等現場編譯成二進制代碼再執行，相比解釋執行的方式，運行效率要高得多。尤其是對于計算密集型查詢、或頻繁重復使用的計算過程，運用代碼生成技術能達到數十倍的性能提升。
代碼生成
很多大數據產品都將代碼生成技術作為賣點，然而事實上往往談論的不是一件事情。比如，之前就有人提問：Spark 1.x 就已經有代碼生成技術，為什么 Spark 2.0 又把代碼生成吹了一番？其中的原因在于，雖然都是代碼生成，但是各個產品生成代碼的粒度是不同的：
o 最簡單的，例如 Spark 1.4，使用代碼生成技術加速表達式計算；
o Spark 2.0 支持將同一個 Stage 的多個算子組合編譯成一段二進制；
o 支持將自定義函數、存儲過程等編譯成一段二進制，如 SQL Server。

本節主要講上面最簡單的表達式編譯。通過一個簡單的例子，初步了解代碼生成的流程。
解析執行的缺陷
在講代碼生成前，回顧一下解釋執行。以上面圖中的表達式 X×5+log(10)X×5+log?(10) 為例，計算過程是一個深度優先搜索（DFS）的過程：
1） 調用根節點 + 的 visit() 函數：分別調用左、右子節點的 visit() 再相加；
2） 調用乘法節點 * 的 visit() 函數：分別調用左、右子節點的 visit() 再相乘；
3）調用變量節點 X 的 visit() 函數：從環境中讀取 XX 的值以及類型。
（……略）最終，DFS 回到根節點，得到最終結果。
@Override public Object visitPlus(CalculatorParser.PlusContext ctx) {
Object left = visit(ctx.plusOrMinus());
Object right = visit(ctx.multOrDiv());
if (left instanceof Long && right instanceof Long) {
return (Long) left + (Long) right;
} else if (left instanceof Long && right instanceof Double) {
return (Long) left + (Double) right;
} else if (left instanceof Double && right instanceof Long) {
return (Double) left + (Long) right;
} else if (left instanceof Double && right instanceof Double) {
return (Double) left + (Double) right;
}
throw new IllegalArgumentException();
}
上述過程中有幾個顯而易見的性能問題：
o 涉及到大量的虛函數調用、即函數綁定的過程，如 visit() 函數，虛函數調用是一個非確定性的跳轉指令， CPU 無法做預測分支，導致打斷 CPU 流水線；
o 在計算前不能確定類型，各個算子的實現中會出現很多動態類型判斷，例如，如果 + 左邊是 DECIMAL 類型，右邊是 DOUBLE，需要先把左邊轉換成 DOUBLE 再相加；
o 遞歸中的函數調用打斷了計算過程，不僅調用本身需要額外的指令，而且函數調用傳參是通過棧完成的，不能很好的利用寄存器（這一點在現代的編譯器和硬件體系中已經有所緩解，但顯然比不上連續的計算指令）。
代碼生成基本過程
代碼生成執行，顧名思義，最核心的部分是生成出需要的執行代碼。
拜編譯器所賜，不需要寫難懂的匯編或字節碼。在 native 程序中，通常用 LLVM 的中間語言（IR）作為生成代碼的語言。JVM 上更簡單，因為 Java 編譯本身很快，利用運行在 JVM 上的輕量級編譯器 janino，可以直接生成 Java 代碼。
無論是 LLVM IR 還是 Java 都是靜態類型的語言，在生成的代碼中再去判斷類型，顯然不是個明智的選擇。通常的做法是在編譯之前就確定所有值的類型。幸運的是，表達式和 SQL 執行調度，都可以事先做類型推導。
所以，代碼生成往往是個 2-pass 的過程：先做類型推導，再做真正的代碼生成。第一步中，類型推導的同時，其實也是在檢查表達式是否合法，很多地方也稱之為驗證（Validate）。
在代碼生成完成后，調用編譯器編譯，得到了所需的函數（類），調用即可得到計算結果。如果函數包含參數，如上面例子中的 X，每次計算可以傳入不同的參數，編譯一次、計算多次。
以下的代碼實現都可以在 GitHub 項目 fuyufjh/calculator 找到。
驗證（Validate）
為了盡可能簡單，例子中僅涉及兩種類型：Long 和 Double

這一步中，將合法的表達式 AST 轉換成 Algebra Node，這是一個遞歸語法樹的過程，下面是一個例子（由于 Plus 接收 Long/Double 的任意類型組合，沒有做類型檢查）：
@Override public AlgebraNode visitPlus(CalculatorParser.PlusContext ctx) {
return new PlusNode(visit(ctx.plusOrMinus()), visit(ctx.multOrDiv()));
}
AlgebraNode 接口定義如下：
public interface AlgebraNode {
DataType getType(); // Validate 和 CodeGen 都會用到
String generateCode(); // CodeGen 使用
List getInputs();
}
實現類大致與 AST 的中的節點相對應，如下圖。

對于加法，類型推導的過程很簡單——如果兩個操作數都是 Long，結果為 Long，否則為 Double。
@Override public DataType getType() {
if (dataType == null) {
dataType = inferTypeFromInputs();
}
return dataType;
}

private DataType inferTypeFromInputs() {
for (AlgebraNode input : getInputs()) {
if (input.getType() == DataType.DOUBLE) {
return DataType.DOUBLE;
}
}
return DataType.LONG;
}
生成代碼
依舊以加法為例，利用上面實現的 getType()，可以確定輸入、輸出的類型，生成出強類型的代碼：
@Override public String generateCode() {
if (getLeft().getType() == DataType.DOUBLE && getRight().getType() == DataType.DOUBLE) {
return “(” + getLeft().generateCode() + " + " + getRight().generateCode() + “)”;
} else if (getLeft().getType() == DataType.DOUBLE && getRight().getType() == DataType.LONG) {
return “(” + getLeft().generateCode() + " + (double)" + getRight().generateCode() + “)”;
} else if (getLeft().getType() == DataType.LONG && getRight().getType() == DataType.DOUBLE) {
return “((double)” + getLeft().generateCode() + " + " + getRight().generateCode() + “)”;
} else if (getLeft().getType() == DataType.LONG && getRight().getType() == DataType.LONG) {
return “(” + getLeft().generateCode() + " + " + getRight().generateCode() + “)”;
}
throw new IllegalStateException();
}
注意，目前代碼還是以 String 形式存在的，遞歸調用的過程中通過字符串拼接，一步步拼成完整的表達式函數。
以表達式 a + 2*3 - 2/x + log(x+1) 為例，最終生成的代碼如下：
(((double)(a + (2 * 3)) - ((double)2 / x)) + java.lang.Math.log((x + (double)1)))
其中，a、x 都是未知數，但類型是已經確定的，分別是 Long 型和 Double 型。
編譯器編譯
Janino 是一個流行的輕量級 Java 編譯器，與常用的 javac 相比，最大的優勢是：可以在 JVM 上直接調用，直接在進程內存中運行編譯，速度很快。
上述代碼僅僅是一個表達式、不是完整的 Java 代碼，但 janino 提供了方便的 API，能直接編譯表達式：
ExpressionEvaluator evaluator = new ExpressionEvaluator();
evaluator.setParameters(parameterNames, parameterTypes); // 輸入參數名及類型
evaluator.setExpressionType(rootNode.getType() == DataType.DOUBLE ? double.class : long.class); // 輸出類型
evaluator.cook(code); // 編譯代碼
實際上，也可以手工拼接出如下的類代碼，交給 janino 編譯，效果是完全相同的：
class MyGeneratedClass {
public double calculate(long a, double x) {
return (((double)(a + (2 * 3)) - ((double)2 / x)) + java.lang.Math.log((x + (double)1)));
}
}
最后，依次輸入所有參數，即可調用剛剛編譯的函數：
Object result = evaluator.evaluate(parameterValues);

References
o Apache Spark - GitHub
o Janino by janino-compiler
o fuyufjh/calculator: A simple calculator to demonstrate code gen technology
2. 查詢編譯執行
代碼生成（Code Generation）技術廣泛應用于現代的數據系統中。代碼生成是將用戶輸入的表達式、查詢、存儲過程等現場，編譯成二進制代碼再執行，相比解釋執行的方式，運行效率要高得多。
上一節表達式編譯中提到，雖然表面上都叫“代碼生成”，但是實際可以分出幾種粒度的實現方式，如表達式的代碼生成、查詢的代碼生成、存儲過程的代碼生成等。本節要講的是查詢級別的代碼生成，有時也稱作算子間（intra-operator）級別，這也是主流數據系統所用的編譯執行方式。
主要參考了 HyPer 團隊發表在 VLDB’11 的文章 。
https://www.vldb.org/pvldb/vol4/p539-neumann.pdf
Volcano 經典執行模型
為什么要用編譯執行？編譯執行有哪幾種實現？
主角是查詢（Query）的編譯執行，看看經典 Volcano 模型是怎么做的。Volcano 模型十分簡單（這也是流行的主要原因）：每個算子需要實現一個 next() 接口，意為返回下一個 Tuple。

Query 1 是一個很簡單的查詢，Project 會調用 Filter 的 next() 獲得數據，Filter 的 next() 又會調用 TableScan 的 next()，TableScan 讀出表中的一行數據并返回。如此往復，直到數據全部處理完。
Query 2 復雜一些，包含一個 HashJoin。HashJoin 的兩個子節點是不對稱的，一邊稱為 build-side，另一邊稱為 probe 或 stream-side。執行時，必須等待 build-side 處理完全部數據、構建出哈希表之后，才能運行 stream-side。
因為這個原因，執行的過程分成了兩個階段（圖中淺灰色的背景）。在 Volcano 模型中，這也很容易實現，試著寫一下 HashJoin 的偽代碼：
Row HashJoin::next() {
// Stage 1: Build Hash Table (HT)
if (HT is not built yet) // 注意：Build 僅在第一次調用 next() 時發生
while ((r = left.next()) != END)
ht.put(buildKey?, buildValue?)
// Stage 2: Probe tuples one by one
while (r = right.next())
if (HT contains r)
output joined row;
}
這個構建哈希表的過程，稱為物化（Materialize），意味著 Tuple 不能繼續往上傳遞，暫存到某個 buffer 里。大多數時候，如執行 Filter 等算子時，Tuple 一路傳上去，稱為 Pipeline。顯然物化的代價是比較高的，希望盡可能多的 Pipeline 避免物化。
Query 3 中的 Aggregate 算子，有類似的情況：在 Aggregate 返回第一條結果前，要把下面所有的數據都聚合完成才行。
稱 HashJoin、HashAgg 這種打斷 Pipeline 的算子為 Pipeline Breaker，使得執行過程分成了不止一個階段。分成多個階段，因為 HashJoin 或 HashAgg 算法本身決定的，跟 Volcano 執行模型無關。
Volcano 的性能問題
Volcano 執行模型勝在簡單易懂，在那個硬盤速度跟不上 CPU 的時代，性能方面不需要考慮太多。然而隨著硬件的進步，IO 很多時候已經不再是瓶頸，這時候人們就開始重新審視 Volcano 模型，產生了兩種改進思路：
1）將 Volcano 迭代模型和向量化模型結合，每次返回一批而不是一個 Tuple；
2）利用代碼生成技術，消除迭代計算的性能損耗。
關于這兩個方案哪個更優，這里有一篇非常棒的論文做了很詳盡的實驗和分析。
http://www.vldb.org/pvldb/vol11/p2209-kersten.pdf
就像表達式解析執行一樣，Volcano 其實是對算子樹的解釋執行，同樣存在這些問題：
o 每產生一條結果就要做很多次虛函數調用，消耗了大量的 CPU 時間；
o 過多的函數調用導致不能很好的利用寄存器。
如果讓去把 Query 1 寫成代碼來執行，會是什么樣的呢？答案非常短，短的令人驚訝：

右圖中用不同顏色標出了原來的算子，其中 condition = true 是一個表達式，按照上一節講解的方法就能生成出代碼，放到這邊 if 的條件上即可。
這兩個的執行效率應該很容易看出差距！生成出的代碼完全消除了虛函數調用，Tuple 幾乎一直在高速緩存甚至寄存器中。論文中也提到，隨便找個本科生手寫代碼，執行性能都能甩迭代模型幾條街。
再看個更復雜的例子找找感覺，以下查詢（記作 Query 4）混合了 Join、Aggragate 甚至子查詢，這些算子是 Pipeline Breaker，執行過程不可避免的分成幾個階段；除此以外，希望其它部分盡可能地做到 Pipeline 執行。

這個例子有點長，相信對代碼生成已經有了些直覺上的理解，這對理解掌握本節的內容大有幫助。
圖中用不同顏色出了 HashJoin、HashAgg 三個算子各自的代碼，可以看出，各自的代碼邏輯被“分散”到了不止一處地方，甚至代碼中已經很難分辨出各個算子，全都融合（Fusion）到一塊。
這就是想要的結果！如何自動生成出這樣的代碼呢？
很多人有個錯覺，以為數據庫查詢過程那么復雜，生成的代碼一定也很復雜吧。其實不然，查詢中復雜的部分，如 HashJoin 中哈希表實現、TableScan 讀取數據的實現等，這些不用生成很多代碼，僅僅只是調用現有的函數即可，如 LLVM IR 可以調用已存在的任何函數。
換個角度看，生成的代碼不過是把這些算子的實現，更高效的方式串聯：算子自身邏輯就像齒輪，生成的代碼好比連接齒輪的鏈條。
HyPer 的解決方案
代碼生成是個純粹的工程問題。工程問題沒有什么不能解的，難就難在找到其中最漂亮的解。如現在這個問題，為了編程的優雅，希望造一個可擴展的框架：不論哪個算子，只要實現某種接口（就像 Volcano 模型要求實現 next() 接口一樣），就能參與到代碼生成中。
模型要求所有算子實現以下兩個接口函數：
o produce()
o consume(attributes, node)
代碼生成的過程總是從調用根節點的 produce() 開始；consume() 類似于一個回調函數，當下層的算子完成自己的使命之后，調用上層的 consume() 來消費剛剛產生的 tuples——注意這里并不是真的消費。
用例子來說明。下面是一個偽代碼版本的若干算子實現。produce() 和 consume() 返回的類型都是生成的代碼片段，這里為了方便演示直接用字符串表示。真實世界中當然要更復雜一些。

表中紅色的字符串是生成的代碼，黑色的則是 code-gen 本身的代碼。回憶一下：代碼生成其實就是用各種手段拼出代碼（字符串）來，沒什么神秘的。

不滿足于偽代碼，可以嘗試閱讀 HyPer 的 論文（生成 LLVM IR），或者 Spark SQL 中的 CodeGenerator 實現（生成 Java 代碼），后者的代碼相對更容易理解些。
思考：這是唯一的解法嗎？
為什么是 produce/consume 呢？是否存在更簡單的解呢？這里給出推導思路。
首先，如果只有一個接口函數，不妨叫 produce()，一定是不夠用的。為什么這么說呢？一個函數充其量只能做出類似 DFS 的效果：每個算子只會被經過一次。這對 Query 1 還不是問題，但對于上文中復雜的 Query 4，HashJoin 的兩部分代碼離得那么遠，用 DFS 就很難做到了。
為了處理 HashJoin，該增加一個怎樣的函數呢？應該類似于一個回調，如 Query 4 中，當 DFS 進行到 ?a=b?a=b 時，希望通過一種某種方式告訴下面的 σx=7σx=7：當拿到結果后，只要用傳給方法去消費這些 Tuples（生成消費這些 Tuples 的代碼）。這個方法，不妨叫做 consume()。
順理成章的，consume() 至少有個參數來傳遞需要消費的 tuples 有哪些列。另外，需要一個參數用來指示：調用者是左孩子還是右孩子？等價于傳 this。
論文提出的 produce/consume 模式可能是唯一正確的方法，即使存在其它算法，猜想也是大同小異。
References

1. Efficiently Compiling Efficient Query Plans for Modern Hardware - VLDB’11
2. SPARK-12795 - Whole stage codegen
3. Everything You Always Wanted to Know About Compiled and Vectorized Queries But Were Afraid to Ask - VLDB’18
   參考鏈接：
   https://ericfu.me/code-gen-of-expression/
   http://ericfu.me/code-gen-of-query/

總結

以上是生活随笔為你收集整理的JIT Code Generation代码生成的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。