Htto的小站
0%

llvm-mem2reg

发表于 更新于 分类于 ，

1. Mem2Reg

mem2reg 主要将 llvm 中的 alloc, load, store 消除，将 LLVM IR 转换为纯的 SSA 形式IR，例如，对于如下代码：

1
2
3
4
5
6
int foo(int x, int y) {
	if(x < 10)
		return 3*x;
	else 
		return 4*y;
}

执行：clang -S -c -Xclang -disable-O0-optnone -fno-discard-value-names -emit-llvm case1.c -o case1.ll 可以得到如下 IR

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
define dso_local signext i32 @foo(i32 noundef signext %x, i32 noundef signext %y) #0 {
entry:
  %retval = alloca i32, align 4
  %x.addr = alloca i32, align 4
  %y.addr = alloca i32, align 4
  store i32 %x, ptr %x.addr, align 4
  store i32 %y, ptr %y.addr, align 4
  %0 = load i32, ptr %x.addr, align 4
  %cmp = icmp slt i32 %0, 10
  br i1 %cmp, label %if.then, label %if.else

if.then:                                          ; preds = %entry
  %1 = load i32, ptr %x.addr, align 4
  %mul = mul nsw i32 3, %1
  store i32 %mul, ptr %retval, align 4
  br label %return

if.else:                                          ; preds = %entry
  %2 = load i32, ptr %y.addr, align 4
  %mul1 = mul nsw i32 4, %2
  store i32 %mul1, ptr %retval, align 4
  br label %return

return:                                           ; preds = %if.else, %if.then
  %3 = load i32, ptr %retval, align 4
  ret i32 %3
}

使用 opt -passes=mem2reg case1.ll -o case1.opt.ll 将改IR转换为 SSA 形式

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
define dso_local signext i32 @foo(i32 noundef signext %x, i32 noundef signext %y) #0 {
entry:
  %cmp = icmp slt i32 %x, 10
  br i1 %cmp, label %if.then, label %if.else

if.then:                                          ; preds = %entry
  %mul = mul nsw i32 3, %x
  br label %return

if.else:                                          ; preds = %entry
  %mul1 = mul nsw i32 4, %y
  br label %return

return:                                           ; preds = %if.else, %if.then
  %retval.0 = phi i32 [ %mul, %if.then ], [ %mul1, %if.else ]
  ret i32 %retval.0
}

可以看到 load,store,alloc 命令全部消除了。

LLVM 实现 mem2reg 的 pass 主要是 PromotePass（/llvm-project/llvm/lib/Transforms/Utils/MemoryOpRemark.cpp），分析这个 Pass 的 run，可以发现，首先它收集了所有的 AllocaInst，然后调用 PromoteMemToReg将 Alloc和相关的 Store，Load替换掉，直到 Alloca 集合收敛到空，这个 Pass 在执行 PromoteMemToReg 前还做了支配树分析和AssumptionCache，支配树分析很容易理解，因为后面计算 IDF 的时候需要用到支配树， AssumptionCache 这里不做讨论。

1.1 PromoteMem2Reg

1.1.1 LLVM 基本优化

路径：llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

1
void PromoteMem2Reg::run();

这是 PromoteMem2Reg 类的方法，这个类接受了 Allocas 和支配树，然后对他们进行处理。这里我们不关注 Debug 的一些细节。

首先它遍历每个 AllocaInst，检测这个 AllocaInst 是否可以被优化

  • 比如，如果这个 AllocaInst 没有被使用，直接删除即可
  • 如果这个 AllocaInst 只在一个 BasicBlock 中被定义（store 命令只出现在一个 BB），那么可以遍历支配树，替换支配节点的所有 load 为 store的值就行了，具体见 rewriteSingleStoreAlloca 这个函数。
  • 如果 Def 和 Use 都在一个 BasicBlock，执行一趟扫描，见 promoteSingleBlockAlloca

后面就是针对一般情况，一个 AllocaInst 在多个 BasicBlock 被 store，load。

首先给每个 BasicBlock 编号，然后计算哪些 BasicBlock 是存活的，即可以被插入 Phi Node的（见ComputeLiveInBlocks)

  • 如果一个BasicBlock 在使用 AI 之前就 Store了，那么这个 BasicBlock 前就不可能被插入 Phi，所以可以认为它不活跃。

1.1.2 DF 计算算法

下面就是 IDF 的计算了，IDF 的计算大体遵循论文中的方法 (A Linear Time Algorithm for Placing-Phi-Nodes 1995) 1. 首先定义 D Edge 和 J Edge * D Edge: Domniator Edge，支配边，支配树上的边 * J Edge 就是 CFG 上无支配关系的边。有个性质就是 对于 J edge 的一条边 \(u\rightarrow v,level_u \ge level_v\) ，简单来说就是不能往下跳，不然\(u\) 的父节点就不能支配 \(u\) 了。 2. 给每个节点分配一个 level，即该节点在支配树上的深度，如果要计算一个节点 \(u\) 的支配边界，可以想到对于任意一个节点 \(v\in \mathrm{SubDomTree}(u)\) ，如果一个J Edge， \(v\rightarrow z\) 满足 \(level_z \le level_u\) 那么 \(z \in DF(u)\) ，根据DF 的定义，借助DJ Graph 很显然有这样的结论。

1.1.3 IDF 计算

由于我们的 Def 即Store 指令可能存在于多个 BasicBlock，所以要计算一个集合的 \(DF\)，即\(DF(S)=\bigcup\limits_{b\in S}DF(b)\)，但是这些 \(DF(S)\) 里的节点还可能继续传播我们的 \(DF\)，所以我们需要引入 \(\mathrm{IDF}\)，即不停的计算 \(IDF(S)=DF(S\cup DF(S\cup DF(S...))\) 直到达到一个不动点，我们可以用不动点算法，但是这样时间复杂度太高了。于是这个论文给出了一些优化方法

考虑一个节点 \(u\) ，\(v \in SubDomTree(u)\) ，如果我们算出了 \(DF(v)\) 那么计算 \(DF(u)\) 遍历 \(u\) 子树遇到节点 \(v\) 可以立即退出，因为 \(v\) 即其子节点的 \(DF\) 贡献一定是 \(DF(v)\) 的子集，对于 \(z \in DF(v)\)，要满足 \(level_z \le v\) 的难度一定小于 \(level_z \le level_u\)，因为 \(level_u \lt level_v\)

那我们 \(O(n)\) 的算法就非常简单了，用一个优先队列，每次优先计算 level 大的节点（当然原文用的 PiggyBank 数据结构，总之思想就是 level 从大到小计算）。下面我们看 LLVM 是如何实现的。

1.1.4 LLVM 实现思路

见：llvm/include/llvm/Support/GenericIteratedDominanceFrontier.h

1
2
template <class NodeTy, bool IsPostDom>
void IDFCalculatorBase<NodeTy, IsPostDom>::calculate(SmallVectorImpl<NodeTy *> &IDFBlocks)

首先，LLVM 先将所有定义的点加入到有个优先队列里面

1
2
3
4
5
6
7
8
9
10
11
12
// PQ 是个优先队列，而且它的比较对象是一个 pair<int,int> ，first 是 level，second 是 DFS 序
using DomTreeNodePair =
      std::pair<DomTreeNodeBase<NodeTy> *, std::pair<unsigned, unsigned>>;
using IDFPriorityQueue =
      std::priority_queue<DomTreeNodePair, SmallVector<DomTreeNodePair, 32>,less_second>;

IDFPriorityQueue PQ;
for (NodeTy *BB : *DefBlocks)
    if (DomTreeNodeBase<NodeTy> *Node = DT.getNode(BB)) {
      PQ.push({Node, std::make_pair(Node->getLevel(), Node->getDFSNumIn())});
      VisitedWorklist.insert(Node);
    }

后面借助优先队列实现的 worklist 算法，每拿出一个节点，我们要考虑它的子树上的孩子，所以这里面还套了一个worklist 不过很好理解，由于每次我们加入的元素 level 一定小于当前的 level，每个节点的访问次数不超过 \(2\) 次（一次是优先队列访问，一次是父节点扫描孩子时访问），所以最终时间复杂度是 \(O(n \;\log n)\), 后面的 \(\log\) 取决于使用的数据结构。

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
while (!PQ.empty()) {
  DomTreeNodePair RootPair = PQ.top();
  PQ.pop();
  DomTreeNodeBase<NodeTy> *Root = RootPair.first;
  unsigned RootLevel = RootPair.second.first;

  // Walk all dominator tree children of Root, inspecting their CFG edges with
  // targets elsewhere on the dominator tree. Only targets whose level is at
  // most Root's level are added to the iterated dominance frontier of the
  // definition set.

  assert(Worklist.empty());
  Worklist.push_back(Root);

  while (!Worklist.empty()) {
    DomTreeNodeBase<NodeTy> *Node = Worklist.pop_back_val();
    NodeTy *BB = Node->getBlock();
    // Succ is the successor in the direction we are calculating IDF, so it is
    // successor for IDF, and predecessor for Reverse IDF.
    auto DoWork = [&](NodeTy *Succ) {
      DomTreeNodeBase<NodeTy> *SuccNode = DT.getNode(Succ);

      const unsigned SuccLevel = SuccNode->getLevel();
      if (SuccLevel > RootLevel)
        return;

      if (!VisitedPQ.insert(SuccNode).second)
        return;

      NodeTy *SuccBB = SuccNode->getBlock();
      if (useLiveIn && !LiveInBlocks->count(SuccBB))
        return;

      IDFBlocks.emplace_back(SuccBB);
      if (!DefBlocks->count(SuccBB))
        PQ.push(std::make_pair(
            SuccNode, std::make_pair(SuccLevel, SuccNode->getDFSNumIn())));
    };

    for (auto *Succ : ChildrenGetter.get(BB))
      DoWork(Succ);

    for (auto DomChild : *Node) {
      if (VisitedWorklist.insert(DomChild).second)
        Worklist.push_back(DomChild);
    }
  }
}