MLIR与Code Generation
MLIR與Code Generation
MLIR多級(jí)中間表示
MLIR 項(xiàng)目是一種構(gòu)建可重用和可擴(kuò)展編譯器基礎(chǔ)架構(gòu)的新方法。MLIR 旨在解決軟件碎片問(wèn)題，改進(jìn)異構(gòu)硬件的編譯，顯著降低構(gòu)建特定領(lǐng)域編譯器的成本，幫助將現(xiàn)有編譯器連接在一起。
MLIR作用
MLIR 旨在成為一種混合 IR，可以在統(tǒng)一的基礎(chǔ)架構(gòu)中，支持多種不同的需求。例如，包括：
? 表示數(shù)據(jù)流圖的能力（例如在 TensorFlow 中），包括動(dòng)態(tài)shape，user-extensible用戶(hù)可擴(kuò)展的操作生態(tài)系統(tǒng)，TensorFlow 變量等。
? 優(yōu)化和轉(zhuǎn)換通常在此類(lèi)圖上完成（例如在 Grappler 中）。
? 能夠跨內(nèi)核（融合，循環(huán)交換，平鋪等）托管高性能計(jì)算風(fēng)格的循環(huán)優(yōu)化，轉(zhuǎn)換數(shù)據(jù)的內(nèi)存布局。
? 代碼生成“降低”轉(zhuǎn)換，例如 DMA 插入，顯式緩存管理，內(nèi)存平鋪，一維和二維寄存器架構(gòu)的矢量化。
? 能夠表示特定于目標(biāo)的操作，例如特定于加速器的高級(jí)操作。
? 在深度學(xué)習(xí)圖上完成的量化和其它圖轉(zhuǎn)換。
? Polyhedral primitives。
? Hardware Synthesis Tools / HLS。
MLIR 是一種常見(jiàn)的 IR，也支持硬件特定的操作。因此，對(duì)圍繞 MLIR 的基礎(chǔ)架構(gòu)的任何投資（例如，對(duì)其進(jìn)行工作的編譯器通過(guò)），都應(yīng)該產(chǎn)生良好的回報(bào)；許多目標(biāo)可以使用基礎(chǔ)架構(gòu)受益。
MLIR 是一個(gè)強(qiáng)大的表示，但也有非目標(biāo)。不嘗試支持低級(jí)機(jī)器代碼生成算法（如寄存器分配和指令調(diào)度）。更適合較低級(jí)別的優(yōu)化器（例如 LLVM）。此外，不希望 MLIR 成為最終用戶(hù)編寫(xiě)內(nèi)核的源語(yǔ)言（類(lèi)似于 CUDA C++）。另一方面，MLIR 提供了代表任何此類(lèi) DSL集成到生態(tài)系統(tǒng)中的主干。
編譯器基礎(chǔ)架構(gòu)
在構(gòu)建 MLIR 時(shí)，從構(gòu)建其他 IR（LLVM IR，XLA HLO 和 Swift SIL）中獲得的經(jīng)驗(yàn)中受益。MLIR 框架鼓勵(lì)現(xiàn)有的最佳實(shí)踐，例如編寫(xiě)和維護(hù) IR 規(guī)范，構(gòu)建 IR 驗(yàn)證器，提供將 MLIR 文件轉(zhuǎn)儲(chǔ)和解析為文本的能力，使用FileCheck 工具編寫(xiě)大量單元測(cè)試，將基礎(chǔ)架構(gòu)構(gòu)建為一組可以以新方式組合的模塊化庫(kù)。
LLVM 有一些不明顯的設(shè)計(jì)錯(cuò)誤，阻止多線程編譯器，處理 LLVM 模塊中的多個(gè)函數(shù)。MLIR 通過(guò)限制 SSA 范圍，減少 use-def 鏈，用顯式替換cross-function引用，解決這些問(wèn)題 symbol reference。
代碼生成（Code Generation）
代碼生成（Code Generation）技術(shù)廣泛應(yīng)用于現(xiàn)代的數(shù)據(jù)系統(tǒng)中。代碼生成是將用戶(hù)輸入的表達(dá)式，查詢(xún)，存儲(chǔ)過(guò)程等現(xiàn)場(chǎng)編譯成二進(jìn)制代碼再執(zhí)行，相比解釋執(zhí)行的方式，運(yùn)行效率要高得多。尤其是對(duì)于計(jì)算密集型查詢(xún)，或頻繁重復(fù)使用的計(jì)算過(guò)程，運(yùn)用代碼生成技術(shù)能達(dá)到數(shù)十倍的性能提升。
很多大數(shù)據(jù)產(chǎn)品都將代碼生成技術(shù)作為賣(mài)點(diǎn)，然而事實(shí)上，往往談?wù)摰牟皇且患虑椤１热?#xff0c;之前就有人提問(wèn)：Spark 1.x 就已經(jīng)有代碼生成技術(shù)，為什么 Spark 2.0 又把代碼生成吹了一番？其中的原因在于，雖然都是代碼生成，但是各個(gè)產(chǎn)品生成代碼的粒度是不同的：
o 最簡(jiǎn)單的，例如 Spark 1.4，使用代碼生成技術(shù)加速表達(dá)式計(jì)算；
o Spark 2.0 支持將同一個(gè) Stage 的多個(gè)算子組合編譯成一段二進(jìn)制；
o 更有甚者，支持將自定義函數(shù)，存儲(chǔ)過(guò)程等編譯成一段二進(jìn)制，例如 SQL Server。
本文主要講上面最簡(jiǎn)單的表達(dá)式編譯。通過(guò)一個(gè)簡(jiǎn)單的例子，初步了解代碼生成的流程。
解析執(zhí)行的缺陷
在講代碼生成之前，回顧一下解釋執(zhí)行。以上面圖中的表達(dá)式 X×5+log(10)X×5+log?(10) 為例，計(jì)算過(guò)程是一個(gè)深度優(yōu)先搜索（DFS）的過(guò)程：
- 調(diào)用根節(jié)點(diǎn) + 的 visit() 函數(shù)：分別調(diào)用左，右子節(jié)點(diǎn)的 visit() 再相加；
- 調(diào)用乘法節(jié)點(diǎn) * 的 visit() 函數(shù)：分別調(diào)用左，右子節(jié)點(diǎn)的 visit() 再相乘；
- 調(diào)用變量節(jié)點(diǎn) X 的 visit() 函數(shù)：從環(huán)境中讀取 XX 的值以及類(lèi)型。
（……略）最終，DFS 回到根節(jié)點(diǎn)，得到最終結(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();
}
上述過(guò)程中，有幾個(gè)顯而易見(jiàn)的性能問(wèn)題：
o 涉及到大量的虛函數(shù)調(diào)用，即函數(shù)綁定的過(guò)程，例如 visit() 函數(shù)，虛函數(shù)調(diào)用是一個(gè)非確定性的跳轉(zhuǎn)指令， CPU 無(wú)法做預(yù)測(cè)分支，從而導(dǎo)致打斷 CPU 流水線；
o 在計(jì)算之前不能確定類(lèi)型，因而各個(gè)算子的實(shí)現(xiàn)中會(huì)出現(xiàn)很多動(dòng)態(tài)類(lèi)型判斷，例如：如果 + 左邊是 DECIMAL 類(lèi)型，而右邊是 DOUBLE，需要先把左邊轉(zhuǎn)換成 DOUBLE 再相加；
o 遞歸中的函數(shù)調(diào)用打斷了計(jì)算過(guò)程，不僅調(diào)用本身需要額外的指令，而且函數(shù)調(diào)用傳參是通過(guò)棧完成的，不能很好的利用寄存器（這一點(diǎn)在現(xiàn)代的編譯器和硬件體系中，已經(jīng)有所緩解，但顯然比不上連續(xù)的計(jì)算指令）。
代碼生成基本過(guò)程
代碼生成執(zhí)行，顧名思義，最核心的部分是生成出需要的執(zhí)行代碼。
拜編譯器所賜，并不需要寫(xiě)難懂的匯編或字節(jié)碼。在 native 程序中，通常用 LLVM 的中間語(yǔ)言（IR）作為生成代碼的語(yǔ)言。而 JVM 上更簡(jiǎn)單，因?yàn)?Java 編譯本身很快，利用運(yùn)行在 JVM 上的輕量級(jí)編譯器 janino，可以直接生成 Java 代碼。
無(wú)論是 LLVM IR 還是 Java 都是靜態(tài)類(lèi)型的語(yǔ)言，在生成的代碼中再去判斷類(lèi)型顯然不是個(gè)明智的選擇。通常的做法是在編譯之前就確定所有值的類(lèi)型。幸運(yùn)的是，表達(dá)式和 SQL 執(zhí)行計(jì)劃都可以事先做類(lèi)型推導(dǎo)。
所以，綜上所述，代碼生成往往是個(gè) 2-pass 的過(guò)程：先做類(lèi)型推導(dǎo)，再做真正的代碼生成。第一步中，類(lèi)型推導(dǎo)的同時(shí)其實(shí)也是在檢查表達(dá)式是否合法，因此很多地方也稱(chēng)之為驗(yàn)證（Validate）。
在代碼生成完成后，調(diào)用編譯器編譯，得到了所需的函數(shù)（類(lèi)），調(diào)用它即可得到計(jì)算結(jié)果。如果函數(shù)包含參數(shù)，例如上面例子中的 X，每次計(jì)算可以傳入不同的參數(shù)，編譯一次，計(jì)算多次。
以下的代碼實(shí)現(xiàn)都可以在 GitHub 項(xiàng)目 fuyufjh/calculator 找到。
驗(yàn)證（Validate）
為了盡可能簡(jiǎn)單，例子中僅涉及兩種類(lèi)型：Long 和 Double
這一步中，將合法的表達(dá)式 AST 轉(zhuǎn)換成 Algebra Node，這是一個(gè)遞歸語(yǔ)法樹(shù)的過(guò)程，下面是一個(gè)例子（由于 Plus 接收 Long/Double 的任意類(lèi)型組合，所以此處沒(méi)有做類(lèi)型檢查）：
@Override public AlgebraNode visitPlus(CalculatorParser.PlusContext ctx) {
return new PlusNode(visit(ctx.plusOrMinus()), visit(ctx.multOrDiv()));
}
AlgebraNode 接口定義如下：
public interface AlgebraNode {
DataType getType(); // Validate 和 CodeGen 都會(huì)用到
String generateCode(); // CodeGen 使用
List getInputs();
}
實(shí)現(xiàn)類(lèi)大致與 AST 的中的節(jié)點(diǎn)相對(duì)應(yīng)，如下圖。
對(duì)于加法，類(lèi)型推導(dǎo)的過(guò)程很簡(jiǎn)單——如果兩個(gè)操作數(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;
}
生成代碼
依舊以加法為例，利用上面實(shí)現(xiàn)的 getType()，可以確定輸入，輸出的類(lèi)型，生成出強(qiáng)類(lèi)型的代碼：
@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)用的過(guò)程中通過(guò)字符串拼接，一步步拼成完整的表達(dá)式函數(shù)。
以表達(dá)式 a + 2*3 - 2/x + log(x+1) 為例，最終生成的代碼如下：
1 (((double)(a + (2 * 3)) - ((double)2 / x)) + java.lang.Math.log((x + (double)1)))
其中，a，x 都是未知數(shù)，但類(lèi)型是已經(jīng)確定的，分別是 Long 型和 Double 型。
編譯器編譯
Janino 是一個(gè)流行的輕量級(jí) Java 編譯器，與常用的 javac 相比它最大的優(yōu)勢(shì)是：可以在 JVM 上直接調(diào)用，直接在進(jìn)程內(nèi)存中運(yùn)行編譯，速度很快。
上述代碼僅僅是一個(gè)表達(dá)式，并不是完整的 Java 代碼，但 janino 提供了方便的 API 能直接編譯表達(dá)式：
ExpressionEvaluator evaluator = new ExpressionEvaluator();
evaluator.setParameters(parameterNames, parameterTypes); // 輸入?yún)?shù)名及類(lèi)型
evaluator.setExpressionType(rootNode.getType() == DataType.DOUBLE ? double.class : long.class); // 輸出類(lèi)型
evaluator.cook(code); // 編譯代碼
實(shí)際上，你也可以手工拼接出如下的類(lèi)代碼，交給 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
參考鏈接：
https://mlir.llvm.org/
https://ericfu.me/code-gen-of-expression/
總結(jié)
以上是生活随笔為你收集整理的MLIR与Code Generation的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。
- 上一篇： OpenArkCompiler方舟编译
- 下一篇： LLVM一些编程语法语义特性