MLIR与Code Generation

MLIR与Code Generation
MLIR多级中间表示
MLIR 项目是一种构建可重用和可扩展编译器基础架构的新方法。MLIR 旨在解决软件碎片问题，改进异构硬件的编译，显着降低构建特定领域编译器的成本，帮助将现有编译器连接在一起。
MLIR作用
MLIR 旨在成为一种混合 IR，可以在统一的基础架构中，支持多种不同的需求。例如，包括：
• 表示数据流图的能力（例如在 TensorFlow 中），包括动态shape，user-extensible用户可扩展的操作生态系统，TensorFlow 变量等。
• 优化和转换通常在此类图上完成（例如在 Grappler 中）。
• 能够跨内核（融合，循环交换，平铺等）托管高性能计算风格的循环优化，转换数据的内存布局。
• 代码生成“降低”转换，例如 DMA 插入，显式缓存管理，内存平铺，一维和二维寄存器架构的矢量化。
• 能够表示特定于目标的操作，例如特定于加速器的高级操作。
• 在深度学习图上完成的量化和其它图转换。
• Polyhedral primitives。
• Hardware Synthesis Tools / HLS。
MLIR 是一种常见的 IR，也支持硬件特定的操作。因此，对围绕 MLIR 的基础架构的任何投资（例如，对其进行工作的编译器通过），都应该产生良好的回报；许多目标可以使用基础架构受益。
MLIR 是一个强大的表示，但也有非目标。不尝试支持低级机器代码生成算法（如寄存器分配和指令调度）。更适合较低级别的优化器（例如 LLVM）。此外，不希望 MLIR 成为最终用户编写内核的源语言（类似于 CUDA C++）。另一方面，MLIR 提供了代表任何此类 DSL集成到生态系统中的主干。
编译器基础架构
在构建 MLIR 时，从构建其他 IR（LLVM IR，XLA HLO 和 Swift SIL）中获得的经验中受益。MLIR 框架鼓励现有的最佳实践，例如编写和维护 IR 规范，构建 IR 验证器，提供将 MLIR 文件转储和解析为文本的能力，使用FileCheck 工具编写大量单元测试，将基础架构构建为一组可以以新方式组合的模块化库。
LLVM 有一些不明显的设计错误，阻止多线程编译器，处理 LLVM 模块中的多个函数。MLIR 通过限制 SSA 范围，减少 use-def 链，用显式替换cross-function引用，解决这些问题 symbol reference。

 

 

 

 

 

 

 

 代码生成（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<AlgebraNode> 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) 为例，最终生成的代码如下：
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

参考链接：
https://mlir.llvm.org/
https://ericfu.me/code-gen-of-expression/