详解 DeepMind 排序算法( 二 )


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] ldrx3,[x0,16] cmpx2,x3 cselx4,x2,x3,le cselx2,x2,x3,ge cmpx2,x1 cselx3,x2,x1,le cselx2,x2,x1,ge cmpx4,x3 cselx5,x1,x4,gt cselx4,x4,x3,ge stpx5,x4,[x0] strx2,[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] staticinlinevoidSort2( long*a, long*b) { intr = *a < *b; longt = r ? *a : *b; *b = r ? *b : *a; *a = t; } // sorts [a,b,c] assuming [b,c] is already sorted staticinlinevoidPartialSort3( long*a, long*b, long*c) { intr = *c < *a; longt = r ? *c : *a; *c = r ? *a : *c; r = t < *b; *a = r ? *a : *b; *b = r ? *b : t; } // sorts [a,b,c] staticinlinevoidSort3( long*a, long*b, long*c) { Sort2(b, c); PartialSort3(a, b, c); } // sorts [a,b,c,d] staticinlinevoidSort4( 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] staticinlinevoidSort5( 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 的博客文章涵盖了这个主题,但我认为分享一些可移植和可执行的代码可能会很有帮助 。
staticinlinevoidInsertionSort( long*A, longn) { longi, 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; } } voidlongsort( long*A, longn) { longt, p, i, j; switch(n) { case0: return; case1: return; case2: returnSort2(A, A + 1); case3: returnSort3(A, A + 1, A + 2); case4: returnSort4(A, A + 1, A + 2, A + 3); case5: returnSort5(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,我认为他们甚至在堆排序中多加了一个步骤,虽然有点慢,但可以防止恶意代码破坏栈 。


推荐阅读