详解DeepMind排序算法

近日，DeepMind在博客发表了一篇阐述排序算法内核的论文。他们借鉴了构建 AlphaGo 积累的有关深度学习的经验，并将其应用于超优化。这激起了我的兴趣，因为作为一名 C 语言库的作者，我一直在寻找机会来提供最好的工具。从某些方面来说，这是 C 语言库的最终目的。虽然有些功能程序员已经习以为常了，但实际上它们是数十年研究的结晶，反复提炼而成的简单且可移植的代码。DeepMind 的这一发现确实居功至伟，但不幸的是，他们未能解释清楚算法。下面，我们来详细看看他们发布的一段汇编代码，这是一个包含三个元素的数组的排序，我们将伪汇编转换为汇编：

/ move37.S
 .equ P,%rax
 .equ Q,%rcx
 .equ R,%rdx
 .equ S,%rsi
move37: mov (%rdi),P
 mov 8(%rdi),Q
 mov 16(%rdi),R
 mov R,S
 cmp P,R
 cmovg P,R
 cmovl P,S
 cmp S,Q
 cmovg Q,P
 cmovg S,Q
 mov R,(%rdi)
 mov Q,8(%rdi)
 mov P,16(%rdi)
 ret
 .type move37,@function
 .size move37,.-move37
 .globl move37
// deepsort1.c
#include
void move37(long *);
int main() {
 long A[3] = {3, 1, 2};
 move37(A);
 printf("%d %d %d
", A[0], A[1], A[2]);

我将此函数命名为 move37()，因为 DeepMind 在文章中将其与 2016 年 AlphaGo 在与李世石的第二场比赛中的第 37 步进行了比较。这一步让专家们感到震惊，他们都认为 AlphaGo 走错了，但实际上错的是专家们，因为 AlphaGo 最终战胜了对手。下面，我们来运行 DeepMind 的代码：

# run this on the shell
cc -o deepsort1 deepsort1.c move37.S
./deepsort1
2 1 3

在我看来，这个运行结果有错。我提供的数组是 {3, 1, 2} ，但排序结果为 {2, 1, 3}。一定是 DeepMind 在骗我们，因为 2 确实不应该出现在 1 之前。我们来看看他们为开源库 LLVM libcxx （https://reviews.llvm.org/D118029）贡献的代码，看看能否澄清这个问题：

// Ensures that *__x, *__y and *__z are ordered according to the comparator __c,
// under the assumption that *__y and *__z are already ordered.
template
inline _LIBCPP_HIDE_FROM_ABI void __partially_sorted_swap(
 _RandomAccessIterator __x, _RandomAccessIterator __y,
 _RandomAccessIterator __z, _Compare __c) {
 using value_type = typename iterator_traits<_RandomAccessIterator>::value_type;
 bool __r = __c(*__z, *__x);
 value_type __tmp = __r ? *__z : *__x;
 *__z = __r ? *__x : *__z;
 __r = __c(__tmp, *__y);
 *__x = __r ? *__x : *__y;
 *__y = __r ? *__y : __tmp;
}

原来如此。这么说实际上 move37() 并不是排序函数。它是一个排序内核，是函数的 sort3() 的一部分。如果 DeepMind 的论文和博客文章能提到这一点就好了，因为初看之下确实很迷惑。下面是改进版的代码，包括缺少的交换操作。

sort3: mov (%rdi),%rcx
 mov 8(%rdi),%rdx
 mov 16(%rdi),%rsi
 mov %rdx,%rax
 cmp %rdx,%rsi
 cmovl %rsi,%rax
 cmovl %rdx,%rsi
 mov %rcx,%rdx
 cmp %rcx,%rsi
 cmovl %rsi,%rdx
 cmovl %rcx,%rsi
 cmp %rax,%rdx
 cmovge %rax,%rcx
 cmovl %rax,%rdx
 mov %rcx,(%rdi)
 mov %rdx,8(%rdi)
 mov %rsi,16(%rdi)
 ret
 .globl sort3
 .size sort3,.-sort3

为了解释这段代码的重要性，我们来考虑一下这个算法在高层面的应用。在第一次尝试自己解决 sort3() 问题时，我想到了下面这段简短的代码：

// sorts [a,b,c]
 if (a > b) SWAP(a, b);
 if (a > c) SWAP(a, c);
 if (b > c) SWAP(b, c);

回头查看 libcxx，发现他们在做同样的事情。上述代码的问题是编译器无法很好地优化。如果尝试编译上面的代码，你会注意到编译器插入了很多分支指令。这就是 DeepMind 试图改进的地方，他们有更聪明的方法来编写这类代码。然而，这些技巧往往不太容易理解。实际上，我喜欢比较直白的代码，因为比较一下就会发现，这些代码与 DeepMind 最先进的汇编代码的基本思路相同。从根本上说，这个问题的基本思想可以归结为三个比较和交换操作：

mov %rdx,%rax  // create temporary
 cmp %rdx,%rsi  // compare
 cmovl %rsi,%rax  // conditional move
 cmovl %rdx,%rsi  // conditional move
/ repeat thrice

上述代码是预先对网络进行排序的最新技术。下面是 DeepMind 的新发现发挥作用的地方。他们发现上面的 mov 指令有时是不必要的。

sort3: mov (%rdi),%rcx
 mov 8(%rdi),%rdx
 mov 16(%rdi),%rsi
 mov %rdx,%rax
 cmp %rdx,%rsi
 cmovl %rsi,%rax
 cmovl %rdx,%rsi
 mov %rcx,%rdx
 cmp %rcx,%rsi
 cmovl %rsi,%rdx
 cmovl %rcx,%rsi
 mov %rdx,%rcx  // <-- wrekt by AlphaDev
 cmp %rax,%rdx
 cmovge %rax,%rcx
 cmovl %rax,%rdx
 mov %rcx,(%rdi)
 mov %rdx,8(%rdi)
 mov %rsi,16(%rdi)
 ret

尝试运行上面的代码，你会发现无论那行被删除的代码是否存在，运行结果都是 100% 正确的。这行代码看似有用，但实际上什么也没做。因此，几十年来计算机科学都没有注意到这个问题，我并不感到惊讶。到这里，AlphaDev 的工作原理也应该变得更加清晰了。从根本上说，DeepMind 构建了一个人工智能，用于检查汇编代码，并随机删除一些代码，看看代码是否会出问题。我这么说并不是要否定 AlphaDev 的智慧，因为我也做了相同的尝试。

sort3: mov (%rdi),%rcx
 mov 8(%rdi),%rdx
 mov 16(%rdi),%rsi
 mov %rdx,%rax  // can it go?
 cmp %rdx,%rsi
 cmovl %rsi,%rax
 cmovl %rdx,%rsi
 mov %rcx,%rdx  // can it go?
 cmp %rcx,%rsi
 cmovl %rsi,%rdx
 cmovl %rcx,%rsi
 mov %rdx,%rcx  // <-- wrekt by AlphaDev
 cmp %rax,%rdx
 cmovge %rax,%rcx
 cmovl %rax,%rdx
 mov %rcx,(%rdi)
 mov %rdx,8(%rdi)
 mov %rsi,16(%rdi)
 ret

另外，DeepMind 的代码还有一些值得商榷之处。上面的代码中还有两条可以去除的 mov 指令，我们可以使用 ARM64 指令集，针对此类问题生成更精简的代码。可以看到，此处我们不需要任何创建临时变量的指令：

sort4: ldp x1,x2,[x0]
 ldr x3,[x0,16]
 cmp x2,x3
 csel x4,x2,x3,le
 csel x2,x2,x3,ge
 cmp x2,x1
 csel x3,x2,x1,le
 csel x2,x2,x1,ge
 cmp x4,x3
 csel x5,x1,x4,gt
 csel x4,x4,x3,ge
 stp x5,x4,[x0]
 str x2,[x0,16]
 ret

最近 Arm 风靡一时，我想上面的例子证实了他们不负盛名。Arm Limited 也是目前开源领域最乐善好施的公司之一。例如，他们的 MbedTLS 库是我迄今为止见过的最被低估的珍宝之一。在使用这个库的时候，我就想过修改 Arm 的代码以在 x86 硬件上更好地工作。我编写了所有这类的汇编优化，将性能提升到与在 x86 上运行 OpenSSL 相同的水平。MbedTLS 是简单的、可移植的、可修改的 C 代码，因此对于任何想要一个简明易懂的汇编加密库的人来说，这都是个好消息。我将自己的做法告诉了 Arm，虽然他们并不觉得有颠覆性，但仍然给予了友善的鼓励。我希望有一天抽出时间学习 DeepMind 的做法，修改我的上游代码。Arm 的 Optimized Routines 库也非常丰富，该库的双精度数转换在质量上无可挑剔。对于 C 库来说，这个库的帮助特别大，因为几十年来，开源社区一直依靠 Sun Microsystems 于 90 代初编写的数学函数。Arm 找到了改进其中几个函数的方法，例如 pow(x,y)。这可是最基本的数学运算之一，因此可以想象其影响力之大。例如，如果你使用 Arm 的解决方案在 x86 机器上实现 pow(x,y)，那么执行相同操作的速度将比英特尔的原生 x87 指令快 5 倍。很幸运 DeepMind 也加入了这个游戏，我冒昧地将他们的 libcxx diff 转换成了简单易读的 C 代码，这样每个人都能欣赏到它的美丽。这是我希望 DeepMind 的论文和博客文章改进的另一个地方，因为你会在这段代码中发现专家用来让编译器生成无分支 MOVcc 指令的规范技巧。

// sorts [a,b]
static inline void Sort2(long *a, long *b) {
 int r = *a < *b;
 long t = r ? *a : *b;
 *b = r ? *b : *a;
 *a = t;
}
// sorts [a,b,c] assuming [b,c] is already sorted
static inline void PartialSort3(long *a, long *b, long *c) {
 int r = *c < *a;
 long t = r ? *c : *a;
 *c = r ? *a : *c;
 r = t < *b;
 *a = r ? *a : *b;
 *b = r ? *b : t;
}
// sorts [a,b,c]
static inline void Sort3(long *a, long *b, long *c) {
 Sort2(b, c);
 PartialSort3(a, b, c);
}
// sorts [a,b,c,d]
static inline void Sort4(long *a, long *b, long *c, long *d) {
 Sort2(a, c);
 Sort2(b, d);
 Sort2(a, b);
 Sort2(c, d);
 Sort2(b, c);
}
// sorts [a,b,c,d,e]
static inline void Sort5(long *a, long *b, long *c, long *d, long *e) {
 Sort2(a, b);
 Sort2(d, e);
 PartialSort3(c, d, e);
 Sort2(b, e);
 PartialSort3(a, c, d);
 PartialSort3(b, c, d);
}

看到 Sort5() 函数后，我觉得自己对 DeepMind 研究的动机有了更好的理解。在 ARM64 上编译 Sort5() 函数，编译器将生成一个包含 11 个寄存器的函数。你在推导一个数学方程式时，大脑里能否时同时思考 11 个变量？可能不行，这就是我们依赖 PartialSort3 这类内核函数的原因。作为有感知力的生物，人类与猴子并没有太大区别。我们变得更加聪明的主要因素是我们能够解决难题，并将其分解为更小的问题。因此，很高兴看到深度学习应用于增强我们的抽象能力。此外，还有一点值得一提，Sort3() 和 Sort5() 本身就是内核，因为它们的目标就是成为传统排序功能的构建块。DeepMind 的博客文章涵盖了这个主题，但我认为分享一些可移植和可执行的代码可能会很有帮助。

static inline void InsertionSort(long *A, long n) {
 long i, j, t;
 for (i = 1; i < n; i++) {
 t = A[i];
 j = i - 1;
 while (j >= 0 && A[j] > t) {
 A[j + 1] = A[j];
 j = j - 1;
 }
 A[j + 1] = t;
 }
}
void longsort(long *A, long n) {
 long t, p, i, j;
 switch (n) {
 case 0:
 return;
 case 1:
 return;
 case 2:
 return Sort2(A, A + 1);
 case 3:
 return Sort3(A, A + 1, A + 2);
 case 4:
 return Sort4(A, A + 1, A + 2, A + 3);
 case 5:
 return Sort5(A, A + 1, A + 2, A + 3, A + 4);
 default:
 if (n <= 32) {
 InsertionSort(A, n);
 } else {
 for (p = A[n >> 1], i = 0, j = n - 1;; i++, j--) {
 while (A[i] < p) i++;
 while (A[j] > p) j--;
 if (i >= j) break;
 t = A[i];
 A[i] = A[j];
 A[j] = t;
 }
 LongSort(A, i);
 LongSort(A + i, n - i);
 }
 break;
 }
}

上述算法展示了 libcxx 的新功能和改进。基本上就是快速排序，只是在递归到较小的切片时切换到排序内核和插入排序。对于 libcxx，我认为他们甚至在堆排序中多加了一个步骤，虽然有点慢，但可以防止恶意代码破坏栈。此时，可能你最想知道的是，我可以使用这个排序算法吗？这些排序网络内核真的能提高排序速度吗？我认为答案需要视情况而定。如果你只想对 long 进行升序排序，上面的代码将比 C 库提供的标准 qsort() 函数快 2 倍。而且你还不需要内核来处理。到目前为止，我确定的是，在我的个人计算机（搭载了英特尔酷睿 i9-12900KS）上，上述函数对 long 进行排序的速度为每秒 255 兆字节。但是，如果我注释掉排序内核：

void longsort(long *A, long n) {
 long t, p, i, j;
 switch (n) {
 case 0:
 return;
 case 1:
 return;
 /* case 2: */
 /* return Sort2(A, A + 1); */
 /* case 3: */
 /* return Sort3(A, A + 1, A + 2); */
 /* case 4: */
 /* return Sort4(A, A + 1, A + 2, A + 3); */
 /* case 5: */
 /* return Sort5(A, A + 1, A + 2, A + 3, A + 4); */
 default:
 if (n <= 32) {
 InsertionSort(A, n);
 } else {
 for (p = A[n >> 1], i = 0, j = n - 1;; i++, j--) {
 while (A[i] < p) i++;
 while (A[j] > p) j--;
 if (i >= j) break;
 t = A[i];
 A[i] = A[j];
 A[j] = t;
 }
 LongSort(A, i);
 LongSort(A + i, n - i);
 }
 break;
 }
}

longsort() 函数的排序速度可以达到每秒 275 兆字节。我在简化算法后，又实现了 7% 的性能提升。这个函数就是 libc 在加载可执行文件时对 elf 符号表进行排序时使用的函数。long 的好处是，它的长度足以存储一个 int 键值对。能够快速排序映射项是非常有用的技巧。结合朴素快速排序与朴素插入排序是我迄今为止找到的最佳解决方案，因为我必须平衡规模和性能。上面的函数编译后只有 181 字节的 x86-64 机器码。由于 DeepMind 的 sort3() 只有 42 字节，我希望我可以牺牲一些大小来获得性能优势。因为到目前为止我发现的第二佳算法是基数排序，性能可达每秒 400 MB，除了依赖于 malloc() 之外，还需要高达 763 字节的二进制。所以很高兴看到这些内核的表现更好。我并不是说 DeepMind 的想法没有价值。我认为值得注意的是，去年 DeepMind 非常慷慨地公开了矢量化快速排序库（当时还是 Google Brain），并以此实现了永远无法挑战的排序霸权。在我的电脑上，使用 Vqsort 对 long 进行排序的速度为每秒 1155 MB，甚至略胜于 djbsort，后者是开源社区中最受欢迎的库之一，尽管除了 int 之外并没有推广至更多的数据类型。两种实现方式都是通过向量化排序网络来实现的。我认为这就是排序网络技术大放异彩的地方。我想，如果不是 AlphaDev 的智能实体还处于起步阶段，它也会采用这种做法。如果从第一原则开始，仅支持基础指令集就非常困难。我认为再等一等，我们可以期待 AlphaDev 未来会取得伟大的成就，因为它会努力应对更艰巨的挑战。此外，DeepMind 压缩了算法的规模，这一点我也很喜欢，因为这并不常见。极限规模编程是我的爱好之一。我曾在一篇博客文章中发布过一个用于 lambda 演算的虚拟机只有 383 字节，还有一个拥有垃圾收集功能的 lisp 机器，只有 436 字节。我也曾在博客中介绍优化 cosmpolitan c 库大小的技巧。我也喜欢 DeepMind 的母公司谷歌，几周前谷歌授予了我一份开源伙伴的奖金，很高兴看到他们也对压缩代码规模充满了热情。很高兴看到他们用我的库来改进向量化快速排序。我只是希望全世界最佳 long 排序不要因为多加了 24 KB 的二进制编码而变成 C++ 庞然大物。仅升序排序就有 23,000 行汇编代码。我迫切希望有一天看到 AlphaDev 能够对这些代码下手。最后，我喜欢一家人工智能公司建造用机器语言编写代码的机器的想法。为什么他们不喜欢这个想法呢？成为机器是机器的本性。作为一名开发人员，我发现 OpenAI 正在创造的未来缺乏这样的想法，他们建立了一个伟大的大型家长式机器，在零和经济中与地球上的每一位开发者竞争，然后吸引世界各地的租客通过政府监管来控制那台机器。OpenAI 承诺要自动化所有的任务，包括我最喜欢的编程，但他们正在努力创造的未来并不是在朝着这个方向前进。我想要的是能够控制一台能够完成我自己无法完成的事情的机器，比如发现排序内核。这才是真正的进步，我认为我们可以去除的每一行汇编代码都是朝着这个梦想积极迈出的一步。

---