Philis的LLVM魔导书(1):实现一个简单的Function Pass来优化你的IR

中文网络上能找到许多介绍LLVM的文章,但大多是概览类。这个系列的目的是用更加具体的实例来让对编译器知识有一定了解的读者掉入LLVM这个史瓦西半径甚大的坑。

比起内部组织形式晦涩难懂的gcc,优化过程不对用户透明的jvm、.net等虚拟机,LLVM的优化pass提供了一种清晰易懂的优化过程组织方式。每个pass对LLVM来说都是一个小的库,用户可以以即插即用的方式来为自己的编译器添加新的优化过程,也可以自由地在已有的优化pass中选择想要的部分。在这篇文章里,我们将要实现一个简单的pass,Dead Blocks Elimination,来优化我们生成的LLVM IR。

环境设置

首先,我们需要设置build环境,根据你的LLVM版本不同,具体过程也会有些区别。旧版的LLVM使用gmake作为其主要的build工具,后续的版本中逐渐在向cmake迁移。默认读者采用的是笔者写文章时的最新的发布版本8.0.0,请同时参考对应版本文档中官方手册的这一节来配置。

<你的LLVM安装目录>/lib/Transforms/下建立一个名为DeadBlock的目录,进入并新建CMakeLists.txt,将以下代码复制入:

1
2
3
4
5
6
add_llvm_library( DeadBlock MODULE
DeadBlock.cpp

PLUGIN_TOOL
opt
)

随后,退入到上级目录,打开<你的LLVM安装目录>/lib/Transforms/CMakeLists.txt,并添加:

1
add_subdirectory(DeadBlock)

至此,LLVM的build工具就可以正常地读取我们添加的新pass了。

题外话:众所周之,软件行业每年进步最快的永远是版本号,笔者实际用的是14年的远古版本3.5.0(为了编译一个远古项目),还在用Makefile来构建pass。而现在开发中的版本号已经爬到了9.0.0+,趴

Dead Blocks Elimination是什么?

这其实是一个原理非常简单的优化过程,常常出现在优化链的中段。假设在对代码进行了一定优化之后,我们得到了如下的IR:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
; ModuleID = 'test.bc'

define i32 @main() {
b1:
br label %b2

dead: ; No predecessors!
br label %b3

b2: ; preds = %b1
br label %b3

b3: ; preds = %dead, %b2
%b = phi i32 [ 1, %b2 ], [ 2, %dead ]
ret i32 %b
}

注意,由于之前已经进行了一些神秘的优化或者代码本身如此,名为dead的这个base block实际是无法被访问到的!b1中的无条件br永远只会跳转到b2,而不是dead。我们现在要实现的DeadBlock pass正是要消去这样的base block。

不过,并不是直接删掉dead就完事了,在b3之中,我们有一个phi node。为了让原来的代码保持正确,我们还需要调整或者删除受影响的phi node

如果你不了解LLVM IR和phi node的话,可以阅读这篇wiki来了解什么是SSA形式的IR,阅读LLVM的手册来了解LLVM IR的基本语法。

第一铲土

在了解要做什么之后,我们就可以开始动工了。在刚刚建立的DeadBlock目录下添加DeadBlock.cpp,这将是我们编写pass的地方,大部分LLVM优化pass都将代码组织在单个.cpp文件中。和所有C++项目一样,讨厌的第一步总是添加所需的header:

1
2
3
4
5
6
#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/IR/CFG.h"
#include <set>

这个pass的主干将看起来是这个样子:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
using namespace llvm;

namespace {
struct DeadBlock : public FunctionPass
{
static char ID;

DeadBlock() : FunctionPass(ID) {}

virtual bool runOnFunction(llvm::Function& F) override {
// pass的入口
// 下一小节,我们将在这里添加代码
};
};
}

// LLVM会利用pass的地址来为这个id赋值,所以初始值并不重要
char DeadBlock::ID = 0;
// 注册pass,这个pass可能会改变CFG,所以将第三个参数设为true
static RegisterPass<DeadBlock> X("deadblock", "Dead blocks elimination pass", true, false);

LLVM的pass有许多种,本次要实现的是一个FunctionPass,也就是说这个pass将会在每个函数上运行一次。因此,我们的DeadBlock构造体需要继承FunctionPass。在搭建完骨架之后,我们就可以在runOnFunction中逐步添加代码来实现具体的功能了。

造楼

我们需要一些本地变量来暂时存放需要的值:

1
2
3
bool changed = false;
std::set<BasicBlock*> visitedSet;
std::set<BasicBlock*> unreachableSet;
  • changed: 用来指示这个pass是否改变了目标function,如果改变了,则要标记为true。这个值将是我们的返回值
  • visitedSet: 我们将会从这个函数的root block开始,遍历这个root block可能会达到的block,被遍历到的block将会存放到这个set中
  • unreachableSet:在得到visitedSet之后,我们可以将其和这个函数中所有block做比较,如果有不在visitedSet中的block,就将其添加到unreachableSet

了解要做什么之后,就来一起向visitedSetunreachableSet里加内容吧:

1
2
3
4
5
6
7
8
9
10
11
12
// 从EntryBlock开始深度优先遍历整个函数内可以访问的BaseBlock
// 将已被访问过的BaseBlock存放在visitedSet中
for (auto i = df_ext_begin<BasicBlock*,std::set<BasicBlock*>>(&F.getEntryBlock(), visitedSet),
e = df_ext_end<BasicBlock*,std::set<BasicBlock*>>(&F.getEntryBlock(), visitedSet);
i != e; i++);

// 遍历函数内所有BaseBlock,将不在vistitedSet中的BaseBlock添加到unreachableSet中
for (BasicBlock & BB : F) {
if (visitedSet.find(&BB) == visitedSet.end()) {
unreachableSet.insert(&BB);
}
}

虽然只是短短几行代码,LLVM库的强大可见一斑。我们无需自己手动实现深度优先遍历,只需调用DepthFirstIterator.h里的df_ext_begindf_ext_end两个模板,就能轻松遍历整个函数,并将访问过的block添加到visitedSet之中。有了unreachableSet之后,我们就可以判断是否会修改目标函数了:

1
2
3
4
// 标记目标函数是否会被修改
if (!unreachableSet.empty()) {
changed = true;
}

最后只要删除掉不想要的block并返回changed,就大功告成:

1
2
3
4
5
6
7
8
for (BasicBlock* BB : unreachableSet) {
for (auto i = succ_begin(BB); i != succ_end(BB); i++) {
i->removePredecessor(BB);
}
BB->eraseFromParent();
}

return changed;

这时,细心的读者可能会说:“诶诶?!是不是漏掉了什么?刚才说要处理受影响的phi node的呢?”不要着急,这段代码里其实已经处理了这部分啦。removePredecessor()函数会通知该block有predecessor已被移除,随后这个block会检查自己是否有会受到影响的phi node并自动做出修改。怎么样,是不是觉得LLVM的库用起来很方便?

测试结果

下面,就来编译并测试一下我们的成果吧。退回到LLVM安装目录的根目录下,在控制台里输入make来编译刚完成的新pass。如果一切顺利的话,你就能在Debug+Asserts/lib/下看到新生成的DeadBlock.so了。

就来用一开始提到的示例IR作为测试代码好了,用LLVM自带的assembler来编译IR到bytecode:

1
llvm-as test.ll

然后,我们用LLVM自带的opt工具来动态载入新pass并执行。重新进到<你的LLVM安装目录>/lib/Transforms/DeadBlock下:

1
opt -load ../../../Debug+Asserts/lib/DeadBlock.so -deadblock < test.bc > optimized.bc

llvm-dis来看一下优化过的代码吧:

1
llvm-dis optimized.bc
1
2
3
4
5
6
7
8
9
10
11
12
; ModuleID = 'optimized.bc'

define i32 @main() {
b1:
br label %b2

b2: ; preds = %b1
br label %b3

b3: ; preds = %b2
ret i32 1
}

和原来的代码对比一下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
; ModuleID = 'test.bc'

define i32 @main() {
b1:
br label %b2

dead: ; No predecessors!
br label %b3

b2: ; preds = %b1
br label %b3

b3: ; preds = %dead, %b2
%b = phi i32 [ 1, %b2 ], [ 2, %dead ]
ret i32 %b
}

可以看到,dead已经被移除,多余的phi node也被去除,替换成了常量。

完整代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/IR/CFG.h"
#include <set>

using namespace llvm;

namespace {
struct DeadBlock : public FunctionPass
{
static char ID;

DeadBlock() : FunctionPass(ID) {}

virtual bool runOnFunction(llvm::Function& F) override {
bool changed = false;

// visitedSet 用于存放已经被访问过的BaseBlock
// unreachableSet 则在最后用于存放无法被访问到的block
std::set<BasicBlock*> visitedSet;
std::set<BasicBlock*> unreachableSet;

// 从EntryBlock开始深度优先遍历整个函数内可以访问的BaseBlock
// 将已被访问过的BaseBlock存放在visitedSet中
for (auto i = df_ext_begin<BasicBlock*,std::set<BasicBlock*>>(&F.getEntryBlock(), visitedSet),
e = df_ext_end<BasicBlock*,std::set<BasicBlock*>>(&F.getEntryBlock(), visitedSet);
i != e; i++);

// 遍历函数内所有BaseBlock,将不在vistitedSet中的BaseBlock添加到unreachableSet中
for (BasicBlock & BB : F) {
if (visitedSet.find(&BB) == visitedSet.end()) {
unreachableSet.insert(&BB);
}
}

// 标记目标函数是否会被修改
if (!unreachableSet.empty()) {
changed = true;
}

// 遍历unreachableSet,通知其successor移除多余的phi node
for (BasicBlock* BB : unreachableSet) {
for (auto i = succ_begin(BB); i != succ_end(BB); i++) {
i->removePredecessor(BB);
}
BB->eraseFromParent();
}

return changed;
};
};
}

// LLVM会利用pass的地址来为这个id赋值,所以初始值并不重要
char DeadBlock::ID = 0;
// 注册pass,这个pass可能会改变CFG,所以将第三个参数设为true
static RegisterPass<DeadBlock> X("deadblock", "Dead blocks elimination pass", true, false);