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

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

這一步中，將合法的表達式 AST 轉(zhuǎn)換成 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();
}
實現(xiàn)類大致與 AST 的中的節(jié)點相對應(yīng)，如下圖。

對于加法，類型推導(dǎo)的過程很簡單——如果兩個操作數(shù)都是 Long，結(jié)果為 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;
}
生成代碼
依舊以加法為例，利用上面實現(xiàn)的 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 形式存在的，遞歸調(diào)用的過程中通過字符串拼接，一步步拼成完整的表達式函數(shù)。
以表達式 a + 2*3 - 2/x + log(x+1) 為例，最終生成的代碼如下：
(((double)(a + (2 * 3)) - ((double)2 / x)) + java.lang.Math.log((x + (double)1)))
其中，a、x 都是未知數(shù)，但類型是已經(jīng)確定的，分別是 Long 型和 Double 型。
編譯器編譯
Janino 是一個流行的輕量級 Java 編譯器，與常用的 javac 相比，最大的優(yōu)勢是：可以在 JVM 上直接調(diào)用，直接在進程內(nèi)存中運行編譯，速度很快。
上述代碼僅僅是一個表達式、不是完整的 Java 代碼，但 janino 提供了方便的 API，能直接編譯表達式：
ExpressionEvaluator evaluator = new ExpressionEvaluator();
evaluator.setParameters(parameterNames, parameterTypes); // 輸入?yún)?shù)名及類型
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)));
}
}
最后，依次輸入所有參數(shù)，即可調(diào)用剛剛編譯的函數(shù)：
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. 查詢編譯執(zhí)行
代碼生成（Code Generation）技術(shù)廣泛應(yīng)用于現(xiàn)代的數(shù)據(jù)系統(tǒng)中。代碼生成是將用戶輸入的表達式、查詢、存儲過程等現(xiàn)場，編譯成二進制代碼再執(zhí)行，相比解釋執(zhí)行的方式，運行效率要高得多。
上一節(jié)表達式編譯中提到，雖然表面上都叫“代碼生成”，但是實際可以分出幾種粒度的實現(xiàn)方式，如表達式的代碼生成、查詢的代碼生成、存儲過程的代碼生成等。本節(jié)要講的是查詢級別的代碼生成，有時也稱作算子間（intra-operator）級別，這也是主流數(shù)據(jù)系統(tǒng)所用的編譯執(zhí)行方式。
主要參考了 HyPer 團隊發(fā)表在 VLDB’11 的文章 。
https://www.vldb.org/pvldb/vol4/p539-neumann.pdf
Volcano 經(jīng)典執(zhí)行模型
為什么要用編譯執(zhí)行？編譯執(zhí)行有哪幾種實現(xiàn)？
主角是查詢（Query）的編譯執(zhí)行，看看經(jīng)典 Volcano 模型是怎么做的。Volcano 模型十分簡單（這也是流行的主要原因）：每個算子需要實現(xiàn)一個 next() 接口，意為返回下一個 Tuple。

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

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

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

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

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

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。