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）的过程：

调用根节点 + 的 visit() 函数：分别调用左，右子节点的 visit() 再相加；
调用乘法节点 * 的 visit() 函数：分别调用左，右子节点的 visit() 再相乘；
调用变量节点 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) 为例，最终生成的代码如下：
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/

MLIR与Code Generation相关推荐

JIT Code Generation代码生成
JIT Code Generation代码生成一．表达式编译代码生成(Code Generation)技术广泛应用于现代的数据系统中.代码生成是将用户输入的表达式.查询.存储过程等现场编译成二进制 ...
link time code generation
为什么80%的码农都做不了架构师?>>> http://msdn.microsoft.com/en-us/magazine/cc301698.aspx 在link时刻做的代码优 ...
MaxCompute 中的Code Generation技术简介
摘要: 前言在<数据库系统中的Code Generation技术介绍>中,我们简单介绍了一下Code Generation技术及其在大规模OLAP系统,特别是大规模分布式OLAP系统中的 ...
代码生成（Code Generation）表达式编译
代码生成(Code Generation)技术广泛应用于现代的数据系统中.代码生成是将用户输入的表达式.查询.存储过程等现场编译成二进制代码再执行,相比解释执行的方式,运行效率要高得多.尤其是对于计算 ...
JAVA 中的代码生成包 CGLIB （Code Generation Library）
JAVA 中的代码生成包 CGLIB (Code Generation Library) CGLIB 是一个功能强大,高性能的代码生成包.它为没有实现接口的类提供代理,为 JDK 的动态代理提供了很好 ...
asm（Java字节码操控框架）和 CGlib（Code Generation Library）
asm概述 asm 是一个 Java 字节码操控框架. 它能够以二进制形式修改已有类或者动态生成类.ASM 可以直接产生二进制 class 文件,也可以在类被加载入Java 虚拟机之前动态改变类行为. ...
文章翻译|TRANX:A Transition-based Neural Abstract Syntax Parser for Semantic Parsing and Code Generation
Abstract 我们提出了TRANX,一个基于transition的神经语义解析器,将自然语言(NL)映射到formal meaning representation(MRs).TRANX对目标MR ...
中间代码生成（Intermediate Code Generation）
中间代码生成(Intermediate Code Generation) 申明语句的翻译类型表达式申明式语句翻译简单赋值语句的翻译数组引用的翻译控制流语句的翻译控制流语句及其SDT 布尔表 ...
论文阅读：Leveraging Code Generation to Improve Code Retrieval and Summarization via Dual Learning
论文题目与链接: [2002.10198] Leveraging Code Generation to Improve Code Retrieval and Summarization via Dua ...

MLIR与Code Generation

MLIR与Code Generation相关推荐

最新文章

热门文章