for循环优化_从Weld论文看执行器的优化技术

Weld 是一个用于数据计算分析的高性能 Runtime&＃xff08;High-performance runtime for data analytics applications&＃xff09;&＃xff0c;使用 Rust 编写&＃xff0c;可以很容易地集成到各种大数据计算框架中&＃xff0c;比如 Spark SQL、NumPy & Pandas、TensorFlow 等&＃xff0c;带来大幅的性能提升。

除了 Weld 本身的贡献&＃xff0c;论文中提到的各种用于执行阶段的优化技术也很有意思&＃xff0c;其中的大部分都借鉴自关系型数据库或编译器。本文除了介绍 Weld 之外&＃xff0c;也是想对这些技术做个梳理。

本文主要内容来自于 Weld 发表在 VLDB&＃39;18 的论文。

整体架构

之前说到&＃xff0c;Weld 是一个用于数据计算的 Runtime&＃xff0c;它的上层通常是一些计算框架&＃xff0c;例如 Spark SQL、NumPy 等。用户用这些计算框架编写程序&＃xff0c;这些框架将用户需要的计算翻译成 Weld 中间表示&＃xff08;IR&＃xff09;&＃xff0c;然后 Weld 对其进行一系列的优化&＃xff0c;最后生成代码并编译运行。

做个类比&＃xff0c;这就像 LLVM 的工作方式一样&＃xff1a;各种语言的编译前端将高级语言翻译成 LLVM IR&＃xff0c;LLVM 再对 IR 做一系列的优化&＃xff0c;最后再编译成二进制。

虽然都是 IR&＃xff0c;但实际上 Weld IR 和 LLVM IR 有很大不同&＃xff1a;

• Weld IR 是声明式的&＃xff1a;只表达计算流程&＃xff0c;不包含具体的实现。比如下面会提到的 Builder&＃xff0c;上层不需要指定用什么方式构建数组或是哈希表等数据结构&＃xff0c;这些是由 Weld 优化器决定的&＃xff1b;
• Weld IR 是 Lazy 的&＃xff1a;只有当需要输出结果时&＃xff0c;相应的 DAG 计算才会真正开始运行。

上图是 Weld 的整体工作过程&＃xff1a;

1. 上层调用 Weld 的 API 输入需要计算的 IR 程序&＃xff0c;它会被解析成 AST&＃xff1b;
2. 当需要执行时&＃xff0c;相关的函数 IR 会被拼在一起&＃xff0c;方便进行整体优化&＃xff1b;
3. Weld 优化器使用一系列的启发式规则进行优化&＃xff0c;注意结果仍然是 AST&＃xff1b;
4. 最后生成代码并借助 LLVM 编译成二进制。

Weld 主要由两个部分组成&＃xff1a;IR 和 Runtime&＃xff0c;接下来我们依次进行介绍。

Weld IR

Weld IR 支持 int、float 等基本数据类型、struct 类型&＃xff0c;以及两种容器类型&＃xff1a;vec 和 dict&＃xff0c;顾名思义&＃xff0c;分别是&＃xff08;变长&＃xff09;数组和字典。另外还支持他们的各种组合&＃xff0c;就像 JSON 那样。

和数据库的执行器不同&＃xff0c;Weld 不考虑数据拉取之类的问题。它假设输入数据都在内存中以数组形式存在&＃xff0c;例如&＃xff1a; int[100]、struct{int, float}[100]。

Weld IR 的计算都通过 Builder 和 Merger 来完成&＃xff0c;由于 Merger 和 Builder 的接口是一样的&＃xff0c;Weld 论文中并没有把二者区分开来。下面我们统称为 Builder。

Builder 提供两个接口方法&＃xff1a;

• merge(b, v)&＃xff1a;向 Builder b 添加新的元素&＃xff1b;
• result(b)&＃xff1a;拿到 b 的结果&＃xff0c;注意之后不能再添加元素了。

下面是使用 Builder 的例子&＃xff1a;

代码中还有个 for&＃xff0c;它的语法是 for(vector, builders, (builders, index, elem) &＃61;> builders)&＃xff0c;用来并行地对数据做处理——也就是往 Builder 里加元素&＃xff0c;这是 Weld 中唯一的计算方式。

for 还可以同时处理多个 Builder&＃xff0c;这个特性在优化的时候很有用&＃xff0c;可以避免同一个数据扫描多次。

Weld IR 还有些别的特性&＃xff08;比如方便编程的 macro&＃xff09;&＃xff0c;但不是本文的重点&＃xff0c;有兴趣的同学自己看原文吧。

Weld Runtime

当上层输入 IR 并发出开始计算的指令时&＃xff0c;就轮到 Weld Runtime 登场了。在代码生成之前&＃xff0c;Weld Runtime 会对 IR 做优化&＃xff0c;优化可以分为两种&＃xff1a;

1. Rule-Based Optimizer (RBO)&＃xff1a;和我们熟悉的 RBO 优化类似&＃xff0c;是基于规则匹配的优化&＃xff1b;
2. Adaptive Optimizer&＃xff1a;运行时 sample 数据&＃xff0c;然后决定用哪种算法执行&＃xff0c;勉强可以对应 CBO。

为什么不是 CBO&＃xff1f;关系型数据库的 CBO 是需要以统计信息为基础的&＃xff0c;但是 Weld 作为一个通用的 Runtime&＃xff0c;上层框架不一定能提供统计信息&＃xff08;比如 NumPy&＃xff09;。

Weld 应用规则是依次进行的&＃xff0c;每次运行一种优化规则&＃xff0c;称为一个 pass。Pass 之间会进行剪枝&＃xff0c;去掉无用的代码。以下我们逐条看看 Weld 做了哪些优化。

Pipeline

Pipeline 在 OLAP 系统中很常见&＃xff0c;最经典的是 HyPer 团队提出的 consume/produce 代码生成机制&＃xff0c;可以在代码生成时尽可能生成 Pipeline 的代码。

为什么需要 Pipeline&＃xff1f;设想一下使用代码生成、但是不使用 Pipeline 会怎么样&＃xff0c;那么 $R_1$ 和 $sigma_{x&＃61;7}$ 就会分成独立的两步&＃xff0c;$R_1$ (即 TableScan&＃xff09;的结果被物化到内存中&＃xff0c;再进行 $sigma_{x&＃61;7}$&＃xff08;Filter&＃xff09;。

而 Pipeline 的代码省略了中间的物化&＃xff0c;仅仅用了一个 if 就解决了 filter&＃xff0c;这个代价要低得多&＃xff1a;计算 if 表达式时相关数据基本还在寄存器或 Cache 里&＃xff0c;充分利用 Data Locality&＃xff0c;这比去内存取数据快 1&＃xff5e;2 个数量级。

Pipeline 优化规则会在 AST 中匹配这样的模式&＃xff1a;A 的输出就是 B 的输入&＃xff0c;对匹配到的节点应用 pipeline 优化&＃xff0c;下面是一个简单例子&＃xff1a;

Horizontal Fusion

Fusion 意为把两段代码融合成一段更精炼的代码&＃xff0c;刚刚说的 Pipeline 也是一种 Fusion。所谓 Horizontal Fusion 是找出被重复处理的数据&＃xff0c;然后将几次处理合在一起。

例如下面图中的 IR&＃xff0c;v0 原本被 loop over 了两次&＃xff0c;如果把两次循环合成一次&＃xff0c;能尽可能利用 Data Locality&＃xff0c;减少一半的内存读取代价。

硬要说的话&＃xff0c;这个规则与关系代数优化中的 Project Merge 规则最相似。论文中给了一个更好的例子来说明它的用处&＃xff1a;像 Pandas 这类的计算框架&＃xff0c;由于 API 设计一次只能处理一列&＃xff0c;必须借助 Horizontal Fusion 实现一次处理多列。

向量化和 Adaptive 优化

向量化&＃xff08;Vectorization&＃xff09;优化也不是新鲜事&＃xff0c;很多编译器&＃xff08;比如 LLVM&＃xff09;都能自动把循环编译成 SIMD 指令&＃xff0c;JVM 甚至可以在运行时生成 SIMD 代码。

SIMD 全称是单条指令、多个数据&＃xff0c;即用一条指令处理多个数据计算&＃xff0c;比如原本计算 4 个整数加法要用 4 次加法指令&＃xff0c;用了 SIMD 之后只要 1 次。没错&＃xff0c;就这么简单&＃xff01;

在这个 pass 中仅处理简单的、没有条件分支的 for 循环&＃xff0c;如果满足这一条件&＃xff0c;优化器会将被循环的数据从 T 转换成 simd[T]&＃xff0c;最后 code-gen 的时候为其生成 SIMD 代码。

那对于带有条件分支的 for 循环&＃xff0c;能否进行向量化呢&＃xff1f;答案是&＃xff0c;可以&＃xff0c;但是不一定有用。

我们先设想一下&＃xff1a;对于有条件分支的 for 循环&＃xff0c;它向量化之后是什么样的&＃xff1f;SIMD 指令本身是没法处理分支的&＃xff08;compare 这种特别简单的除外&＃xff09;&＃xff0c;如果一定要用 SIMD&＃xff0c;可以假设分支条件全都为 true 或 false&＃xff0c;最后根据条件表达式的计算结果&＃xff08;true 或 false&＃xff09;&＃xff0c;利用 select 指令选出相应的结果。

这种方式相比普通的带分支的指令&＃xff0c;有得有失&＃xff1a;

• 优势&＃xff1a;用 SIMD 指令集可以加速计算&＃xff1b;
• 劣势&＃xff1a;原本只要算一个分支&＃xff0c;现在两个分支都要算。

注&＃xff1a;另一个优势是&＃xff0c;SIMD 去掉了条件跳转&＃xff0c;不存在打断 CPU 流水线的问题。但是论文中没有提到这一点&＃xff0c;我猜测可能是它的影响因素比较小&＃xff0c;或是作者没有找到一个合适的代价计算方式。

论文只给出了对 if(cond, merge(b, body), b) 这样单分支条件的代价建模&＃xff0c;有兴趣的同学可以看原论文上的式子。这里只说一个粗糙的结论&＃xff1a;当选择率&＃xff08;即进入 if-body 的概率&＃xff09;很小时&＃xff0c;有分支的代码更优&＃xff1b;当选择率比较大时&＃xff0c;SIMD 代码更优。

我们之前说过&＃xff0c;Weld 假设上层无法提供统计信息&＃xff0c;因而在这一步&＃xff0c;由于缺乏关键的选择率信息&＃xff0c;它只能采取一种 Adaptive 的思路&＃xff1a;同时生成有分支的代码和 SIMD 代码&＃xff0c;运行时&＃xff0c;首先对输入数据做个 Sample 估算一下选择率&＃xff0c;再决定走哪个算法。

选择率&＃xff08;selectivity&＃xff09;这个概念在数据库优化器中也很常用&＃xff0c;比如估算 Row Count 时就频繁用到了选择率估计。如果能在优化时直接拿到这个信息&＃xff0c;想必不需要这么折腾。

Adaptive Hash Table

Weld 的 dictbuilder 和 groupbuilder 中都需要构建哈希表&＃xff0c;这里也有个 trade-off&＃xff1a;是用 Partitioned Hash Table 还是 Global Hash Table&＃xff1f;

• Partitioned Hash Table 是将 build 过程分成两步&＃xff0c;先各个线程本地做 build&＃xff0c;最后再 merge 成完整的结果&＃xff1b;
• Global Hash Table 只有一张全局的哈希表&＃xff0c;通过加锁等方式做了控制并发写入。

一般而言&＃xff0c;如果 Group by 的基数&＃xff08;Cardinality&＃xff09;比较小&＃xff0c;Partitioned 方式更有优势&＃xff0c;因为并发冲突会很多&＃xff1b;相反&＃xff0c;如果基数很大&＃xff0c;Global 更占优势&＃xff0c;因为无需再做多一次 merge。

Weld 的做法很巧妙地实现了二者取折中&＃xff1a;各线程先写到本地的哈希表&＃xff0c;但如果大小超过阈值&＃xff0c;就写到全局的哈希表。最后把本地数据再 merge 进全局哈希表。这个实现被它称为 Adaptive Hash Table。

Misc.

Weld 中还有还有一些优化手段&＃xff0c;比较简单&＃xff1a;

循环展开&＃xff08;Loop Unrolling&＃xff09;是编译器中很常见的优化&＃xff0c;如果编译期已知 for 循环的次数很小&＃xff08;例如&＃xff0c;对于一个 N*3 的矩阵&＃xff0c;第二维度长度仅为 3&＃xff09;&＃xff0c;就将循环展开&＃xff0c;避免条件跳转指令打断 CPU 流水线。

数组预分配&＃xff08;Preallocation&＃xff09;在矩阵运算中也很有用&＃xff0c;例如&＃xff0c;默认 vecbuiler 的实现是自动生长的动态数组。如果预先知道数组长度&＃xff0c;就能避免数组生长的拷贝代价。

评估和总结

下面是 Weld 官网放出的性能评估&＃xff0c;对于文中提到的这几个框架&＃xff0c;的确做到了可观的性能提升。

注&＃xff1a;这里 TensorFlow 性能是用 CPU 运行的&＃xff0c;而非 GPU。

Weld 的最大贡献是抽象出了一个通用的执行器 Runtime。这个抽象的层级要比“代码生成”中的“代码”&＃xff08;比如 LLVM IR&＃xff09;高级&＃xff08;high-level&＃xff09;不少&＃xff0c;但又比关系代数或是线性代数低级&＃xff08;low-level&＃xff09;&＃xff0c;从而有更好的通用性。更可贵的是&＃xff0c;Weld IR 仅仅包含 Builder 以及 for、if 这些最基本的语句&＃xff0c;极其之简单。

上文提到的很多优化规则&＃xff0c;不少来源于编译器或关系型数据库。例如 Pipeline Fusion 的思想&＃xff0c;在编译器中其实也有体现——编译器会尽可能连续的利用寄存器、避免 store/load。但是 Weld IR 独特的抽象层级令它能做层级更高的优化&＃xff0c;达到和数据库的 Pipeline 一样的效果。

References

1. Evaluating End-to-End Optimization for Data Analytics Applications in Weld (VLDB&＃39;18)
2. Efficiently Compiling Efficient Query Plans for Modern Hardware (VLDB&＃39;11)
3. Weld - Official Website

本文作者&＃xff1a; &＃64;Eric Fu 原文链接&＃xff1a; https://ericfu.me/weld-the-query-exeution-engine/版权声明&＃xff1a; 本文章采用 BY-NC-SA 4.0 许可协议。转载请注明出处&＃xff01;