Chaitin图着色算法

图着色问题

给定一个无向图, 其中是顶点的集合, 是边的集合, 要求为图中的每一个顶点分配一个颜色, 并且相邻的两个顶点不能有同样的颜色. 图着色问题就是在满足上述约束的条件下, 使用最少的颜色为图着色. 这个最少的颜色数量称为图的色数(Chromatic Number), 用符号表示.

但并不打算在这里直接讨论图着色的精确解, 因为这是个NP-Hard问题, 精确解的算法复杂度的指数级的, 即使解出来对于大型图的编译过程也没有实际意义. 接下来我们讨论一个近似解法.

算法简述

假设我们有个寄存器, 个变量. 首先计算出每个变量的活跃区间, 也就是变量从被定义到最后一次被使用的区间. 然后构造一个干扰图(Interference Graph), 图中的每一个顶点代表一个变量, 或者说是一个虚拟寄存器. 图中的每一条边代表变量之间的干扰关系, 也就是如果两个变量的活跃区间有重叠(Overlap), 那么它们之间就有一条边相连. 因此这个图构造出来就是一个无向图. 接下来要为这个变量分配寄存器, 也就是用种颜色为这个图着色, 这就变成了一个图着色问题.

但这个图着色和之前提到的图着色问题有一些区别, 首先我们不要求使用最少的颜色, 只要不超过种颜色就可以了. 其次, 如果图的色数, 那我们可以通过溢出("Spill")到内存的方式将某个变量从图中删除, 然后对新图进行着色. 这个过程可以不断迭代, 直到最终的图的色数小于等于为止.

接下来讨论怎么具体解决这个问题. 关于活跃区间的构建, 这里不做赘述, 可以参考另一篇寄存器分配算法线性扫描. 假设已经有了活跃区间, 接着是构建干扰图.

干扰图可以用位矩阵(Bit Matrix)+邻接向量(Adjacency Vector)来表示. 位矩阵的大小是, 其中如果节点和之间相邻, 那么, 否则. 因此位矩阵一定是一个对称矩阵. 但这个矩阵通常是稀疏的, 处理的时候效率不高. 因此引入邻接向量. 图中每一个节点都有一个邻接向量, 这个邻接向量中储存了所有和这个节点相邻的节点的索引. 在接下来遍历干扰图的时候, 位矩阵+邻接向量的组合可以提高效率.

接下来对这个干扰图进行化简. 容易观察到这样一个事实: 如果一个节点的度小于, 那么这个节点一定可以被着色. 因此如果从图中移除这个节点, 并且更新所有与之相邻的节点的度, 那么这个图的色数不会改变. 于是我们可以先遍历这个干扰图, 找到所有度小于的节点, 将它们从图中移除, 并推入一个栈中.

如果幸运的是, 干扰图化简到最后是个空图, 那么我们可以直接从栈中弹出所有节点并为它们着色(分配寄存器). 但如果最终的图中剩下一些节点(这些节点的度一定都大于或等于), 那么就需要将它们溢出了.

溢出操作就是不断进行尝试, 从原始的图中删除一个或者多个最后剩下的节点, 然后对新图进行着色, 看新图的色数是否小于等于. 如果是, 那么就可以将这些节点溢出到内存中; 如果不是, 那么就调整溢出策略, 重新尝试, 直到最终的图的色数小于等于为止. 最差的情况就是将所有剩下的节点都溢出到内存中, 因此这个问题一定是有解的.

下面是一个示例, 对于这样一段汇编($a和$b只表示两个变量, 并非汇编用法), 寄存器均为虚拟寄存器, 因此数量是无限的:

ld	t0, $a # 加载变量`a`到寄存器
ld	t1, $b # 加载变量`b`到寄存器
add	t2, t0, t1 # 将`a`和`b`相加, 得到变量`c`
add t3, t0, t0 # 将`a`和`a`相加, 得到变量`d`
add t4, t0, t2 # 将`a`和`c`相加, 得到变量`e`

因此可以得到活跃区间表如下

然后构建干扰图, 其中变量a和b,c,d都有重合的活跃区间, 变量c和d也有重合的活跃区间. 因此画出干扰图如下:

假设我们现在只有个物理寄存器, 因此要修改上面的汇编代码, 将变量重映射到物理寄存器或者内存中. 对图进行化简, 显然节点b,e的度都小于, 因此可以将它们从图中移除, 得到

现在节点a,c,d的度都为2, 恰好等于寄存器数量, 因此不能直接移除, 需要选择其中的一个或者多个节点进行溢出. 溢出规则通常会选择溢出成本最小的节点, 计算溢出成本又有不同的方法. 这里就简单选择节点a进行溢出. 更新图, 得到

此时节点3,4的度都为1, 因此将它们从图中移除, 并推入栈中, 最后会得到一个空图和栈:

依次从栈中弹出节点, 并考虑它们在原始图中的邻接关系进行寄存器分配. 首先节点d的邻接节点是a,c, a被溢出, c未分配, 因此将节点d分配到寄存器R1中. 接着节点c的邻接节点是a,d, a被溢出, d已经分配到寄存器R1, 因此将节点c分配到寄存器R2中. 节点b的邻接节点是a, a被溢出, 所以将节点b分配到寄存器R1或者R2都行. 节点e是孤立节点, 因此分配到寄存器R1或者R2都行. 最终得到的寄存器分配如下:

算法细节

对于上面的算法, 还有一些细节需要注意.

节点合并

有一些不必要的寄存器拷贝操作, 比如当变量a和b的活跃区间没有重叠, 但存在指令a=b的时候, 这就是一个不必要的拷贝操作. 如果不做任何处理, 变量a和b可能会被分配到不同寄存器中, 但实际上是没有必要的. 例如:

a = x + 1;
b = a;
b = y + b;

这个时候a和b的活跃区间不重叠, 可以将干扰图上的a和b合并. 合并的方式就是合并a和b的邻接向量. 比如在这里a的邻接向量是x, b的邻接向量是y, 那么合并后新节点ab的邻接向量就是[x,y]. 这个时候可以构建出一个新图. 如果在旧图上直接操作, 后续计算会变得比较麻烦, 因此论文中的操作是在合并后修改IR, 然后基于新的IR重新构建干扰图.

邻接关系的储存方式

除了邻接向量+位矩阵的方式, 还可以用邻接表来储存邻接关系. 如果使用邻接表的话, 那么可以有以下的基本数据结构:

typedef int register_t; // 寄存器类型

typedef struct {
	std::string name; // 节点名称
	int value; // 节点值
	std::vector<Node*> neighbors; // 邻接节点
	register_t reg; // 寄存器分配, -1表示未分配
	bool spilled; // 是否溢出
} Node;

typedef struct {
	std::vector<Node*> nodes; // 图中的所有节点
	std::stack<Node*> alloc_stack; // 分配栈

	void removeNode(Node* node) {
		// 邻接节点列表中删除节点
		// 从图中删除节点
	}
} Graph;

int simplifyGraph(Graph &g) {
	int degree_less_than_R = 0;
	int degree_greater_than_R = 0;
	// iterate the nodes
	// if node.degree < R, g.alloc_stack.push(node); then removeNode(node)

	// if graph is empty (!degree_greater_than_R && !degree_less_than_R)), return 0
	// if graph is not empty, recusively call simplifyGraph(g)
	// if graph cannot be simplified, return -1 (degree_greater_than_R && !degree_less_than_R)
}

void spill(Graph &g) {
	// rank the nodes
	// choose the node to spill

	// increasingly spill nodes according to the rank
	for (auto& node : nodes_rank) {
		g.removeNode(node);
		node.spilled = true; // mark the node as spilled
		if (simplifyGraph(g_tmp) == 0)
			break;
	}
	// or other strategies
	// there must be a solution to make the graph empty
	// the worst case is to spill all nodes
}

void getFreeRegister(Node* node) {
	// get a free register for the node
	// if no free register, return -1
	// else return the free register
}

void allocate(Graph &g) {
	while (!g.alloc_stack.empty()) {
		Node* node = g.alloc_stack.top();
		g.alloc_stack.pop();
		if (!node->spilled) {
			node->reg = getFreeRegister(node); // get a free register
		}
	}
}

if (simplifyGraph(g) != 0)
	spill(g); // spill the graph
allocate(g); // allocate registers

这段代码写得很👎, 它是有问题的, 大概能说明一个思路就行, 不要细究里面的问题🤪.

溢出策略

理论上讲, 溢出策略应该选择代价最小的. 怎么计算这个代价呢? 假设我们对干扰图化简到最后得到了一个非空图, 这个图中有个节点. 最简单的策略就是按照度数从高到低依次尝试溢出, 直到干扰图能被化简为空图. 这简单直接, 符合直觉, 因为度数越高的节点, 溢出后对图的影响越大, 似乎最有可能让干扰图化简为空图. 但显然它不一定是最优的, 如果一个变量在某个区域中是核心变量, 它使用频繁且广泛, 那它的溢出代价就会极高. 比如循环计数器. 它的活跃区间是整个循环, 而且每次循环都要使用它, 因此它的度数非常高, 但显然你不可能溢出它.

那么可以采用一些别的策略, 例如优先溢出那些对引入内存开销最小的节点. 假设每个变量的load/store都要花费1个周期, 可以计算每个变量如果被溢出产生的内存开销, 例如在一个嵌套层级为n的循环中, 某个溢出变量x的内存开销就是2^n. 这种做法的好处是能尽量减少最昂贵的溢出(这样你就不会溢出循环计数器了), 但溢出变量的度数可能比较小, 需要溢出很多个变量才能让图化简为空图. 这会导致一些问题, 一是干扰图迭代次数会增加, 二是所有溢出的变量的内存开销加起来并不一定会降低, 虽然单个溢出变量的内存开销小, 但要溢出的变量多了.

于是有了一种折中的策略, 对每个要尝试溢出的变量, 计算它的度数和内存开销的乘积, 或者进行其它的加权处理. 这个值越大, 表明这个变量的溢出优先级越高. 这种方法应该能在以上两种情况之间取得一些平衡, 但缺点是计算代价是最高的.

寄存器重用(Register Reload)

如果要对某个节点做溢出操作, 就要在定义它时插入store指令, 在每次使用它时插入load指令. 并且这个溢出操作是永久性的, 如果某个节点最终被决定溢出, 那么它被定义后将一直储存在内存中, 不会再为它分配寄存器.

a = x + 1;
c = x * 2;
d = x + 3;

如果x是被溢出的, 那么上面的代码会被转换为

load R1, [address_of_x]
add R2, R1, #1
store R2, [address_of_x]
; ...something else...
load R1, [address_of_x]
mul R3, R1, #2
store R3, [address_of_x]
; ...something else...
load R1, [address_of_x]
add R4, R1, #3
store R4, [address_of_x]

所以如果x被决定全局溢出, 但在某个局部区域中多次使用, 就会因为内存操作导致性能下降.

显然这里存在一个可优化的点, Chaitin算法是一个全局分配算法, 它基于全局中寄存器使用数量的峰值来决定寄存器的分配. 但在某个局部区域中, 很可能寄存器的必须使用量小于寄存器总数. 因此在局部区域中可能找到一个空闲的寄存器分配给x. 如果x的活跃区间横跨了几个局部区域, 那么每个局部区域中都可能找到空闲的寄存器分配给x. 但并不保证每次找到的空闲寄存器都是同一个.

换句话说, 如果使用Chaitin算法+局部优化, 那一个变量被决定溢出, 并不代表它在整个活跃区间都在内存中, 而是大部分时间在内存中, 剩下的时间可能在不同的临时寄存器中.

局部优化在CodeGen阶段进行, 不属于Chaitin核心算法的一部分, 因此在这里不做讨论.

如果寄存器已满

Chaitin算法的分配结果中可能出现这样一个情况, 某个时刻寄存器已经满了, 又要使用溢出的变量进行计算(因为分配寄存器的时候使用了全部个寄存器, 并非个寄存器). 如果在x86上这还说得过去, 因为x86指令允许直接从内存地址读数据进行计算, 但在RISC-V这种"load-store"架构上, ALU只能从寄存器中读数据进行计算, 去哪找寄存器给溢出变量呢?

这个时候就需要选择一个寄存器, 将它其中的未溢出变量临时储存到内存中, 然后将溢出变量加载到这个寄存器中进行计算. 这个被替换掉的变量称为"牺牲品"变量(Victim Variable), 编译器会为它插入load/store指令, 临时将它放到内存中. 当溢出变量计算完成后, 再将牺牲品变量从内存中加载回来.

但关于牺牲品变量如何选择就不是Chaitin算法本身解决的问题了. Chaitin算法是一个全局算法, 它只能决定某个变量在活跃期间的大部分时间内待在哪里, 对局部优化无能为力.

论文原文

链接