DeepMind 新作 AlphaDev ---- 强化学习探索更优排序算法-轻识

前言

DeepMind 最近在 Nature 发表了一篇论文 AlphaDev[2, 3]，一个利用强化学习来探索更优排序算法的AI系统。

AlphaDev 系统直接从 CPU 汇编指令的层面入手去探索更优的排序算法，因为相对于高级编程语言来说，在汇编指令层级对存储和寄存器的操作可以更加的灵活，所以能发现更多潜在的调优策略。

在 AlphaDev 的论文中，只关注探索短序列排序：

定长序列排序（比如 sort3 算法只能对长度为3的序列进行排序）
变长序列排序（比如 variable sort5 算法可以对长度为1~5的变长序列进行排序）

而对于长序列的排序，可以被分解为短序列的排序。

DeepMind 通过 AlphaDev 发现了比目前人工调优算法更优的定长短序列排序算法 sort3，sort4 和 sort5 ，并且已经将代码提交到了 LLVM 标准 C++ 库[4] 。

简单来说，AlphaDev 将探索更高效排序算法的过程，建模为一个单玩家的汇编游戏（single-player game， AssemblyGame）。

游戏的过程就是玩家从 CPU 汇编指令集合中，选取一系列的指令组合得到一个新的排序算法。不过这个过程是非常有挑战的，玩家需要考虑，汇编指令的组合空间并最终得得到一个正确和高效的算法。

该游戏主要包括以下难点：

汇编游戏的搜索空间和围棋类似（10^700）
只要有一条指令没弄对，可能就会导致整个算法错误

AlphaDev 系统详解

将排序算法表示为 CPU 汇编指令

首先来看一个简单的变长（variable sort2）短排序函数的 C 代码实现，排序结果从小到大：

void variable_sort_2(int length, int *a) {
    switch (length) {
    case 0:
    case 1:
        return;
    case 2:
        int tmp = a[0];
        // a[0] 保存两者之间的最小值
        a[0] = (a[1] < a[0]) ? a[1] : a[0];
        // a[1] 保存两者之间的最大值
        a[1] = (a[1] < tmp) ? tmp : a[1];
        return;
    }
}

通过 gcc 生成对应的汇编代码，我用的 gcc 版本是 11.3.0，命令 gcc -S -O1 -o sort2.s sort2.c

汇编代码只保留了核心部分，生成的结果和论文中的示例有些许不同但是原理是一致的：

variable_sort_2:  
.LFB0:
	; %edi 寄存器保存参数 length 的值
	; cmpl 指令对比 %edi 和 常量 2
	cmpl	$2, %edi 
	; 相等就跳转到 .L3 标签处，
        ; 对应 C 代码的 case 2
	je	.L3
.L1:
	; 不等于 2 就直接返回，
        ; 对应 C 代码 case 0 和 1
	ret 
.L3:
	; 将 a[0] 赋值给寄存器 %edx 
	movl	(%rsi), %edx
	; 将 a[1] 赋值给寄存器 %eax 
	movl	4(%rsi), %eax
	; 对比 %edx 和 %eax
	cmpl	%edx, %eax
	; 将 %edx 赋值给 %ecx
	movl	%edx, %ecx
	; cmov 是条件移动指令根据 cmpl 
	; 指令的结果判断是否执行
	; 如果 %eax <= %edx 
	; 则将 %eax 赋值给 %ecx
	cmovle	%eax, %ecx
	; 此时 %ecx 保存了最小值
	; 将 %ecx 赋值给 a[0]
	movl	%ecx, (%rsi)
	; 如果 %eax 小于 %edx
	; 则将 %edx 赋值给 %eax
	cmovl	%edx, %eax
	; 此时 %eax 保存了最大值
	; 将 %eax 赋值给 a[1]
	movl	%eax, 4(%rsi)
	jmp	.L1

一般来说汇编程序所做的事情基本都是，将内存的值复制到寄存器，然后对寄存器的值作修改，再将寄存器的值写回到内存中。

而 AlphaDev 系统只关注 x86 处理器架构所支持的汇编指令集合的一个子集。

每条汇编指令的格式均为：操作码<操作数A, 操作数B> 比如：

mov 移动指令，表示将 A 的值赋值给 B
cmp 比较指令，相当于执行 A - B 操作，但是不会对 A 和 B 做修改，而是根据相减的结果设置特殊的 flag 寄存器，更多内容可以参考[5]
cmovX 条件移动指令，根据 X 和 flag 寄存器的值判断是否执行将 A 赋值给 B 的操作，一般都是出现在 cmp 指令之后。X 可以是 L （是否满足小于条件）， G （是否满足大于条件），LE （是否满足小于或等于条件），GE （是否满足大于等于条件）。
jX 条件跳转指令，根据 X 和 flag 寄存器的值判断是否执行跳转到指定标记位置操作，A 可以是汇编程序代码中的标记位置，如上面所示汇编代码的 .L1 和 .L3。X 可以是 NE （是否不等于），E （是否等于）或者可以填表示无条件跳转。