詳解 DeepMind 排序算法
加入技術交流群
DeepMind 的這一發(fā)現(xiàn)確實居功至偉，但不幸的是，他們未能解釋清楚算法。下面，我們來詳細看看他們發(fā)布的一段匯編代碼，這是一個包含三個元素的數(shù)組的排序，我們將偽匯編轉(zhuǎn)換為匯編：
/ move37.S .equ P,%rax .equ Q,%rcx .equ R,%rdx .equ S,%rsimove37: 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#includevoid move37(long *);int main() { long A[3] = {3, 1, 2}; move37(A); printf("%d %d %d\n", A[0], A[1], A[2]);我將此函數(shù)命名為 move37()，因為 DeepMind 在文章中將其與 2016 年 AlphaGo 在與李世石的第二場比賽中的第 37 步進行了比較。這一步讓專家們感到震驚，他們都認為 AlphaGo 走錯了，但實際上錯的是專家們，因為 AlphaGo 最終戰(zhàn)勝了對手。下面，我們來運行 DeepMind 的代碼：# run this on the shellcc -o deepsort1 deepsort1.c move37.S./deepsort12 1 3在我看來，這個運行結(jié)果有錯。我提供的數(shù)組是 {3, 1, 2} ，但排序結(jié)果為 {2, 1, 3}。一定是 DeepMind 在騙我們，因為 2 確實不應該出現(xiàn)在 1 之前。我們來看看他們?yōu)殚_源庫 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.templateinline _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() 并不是排序函數(shù)。它是一個排序內(nèi)核，是函數(shù)的 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，發(fā)現(xiàn)他們在做同樣的事情。上述代碼的問題是編譯器無法很好地優(yōu)化。如果嘗試編譯上面的代碼，你會注意到編譯器插入了很多分支指令。這就是 DeepMind 試圖改進的地方，他們有更聰明的方法來編寫這類代碼。然而，這些技巧往往不太容易理解。實際上，我喜歡比較直白的代碼，因為比較一下就會發(fā)現(xiàn)，這些代碼與 DeepMind 最先進的匯編代碼的基本思路相同。從根本上說，這個問題的基本思想可以歸結(jié)為三個比較和交換操作：
mov %rdx,%rax // create temporary cmp %rdx,%rsi // compare cmovl %rsi,%rax // conditional move cmovl %rdx,%rsi // conditional move/ repeat thrice上述代碼是預先對網(wǎng)絡進行排序的最新技術。下面是 DeepMind 的新發(fā)現(xiàn)發(fā)揮作用的地方。他們發(fā)現(xiàn)上面的 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嘗試運行上面的代碼，你會發(fā)現(xiàn)無論那行被刪除的代碼是否存在，運行結(jié)果都是 100% 正確的。這行代碼看似有用，但實際上什么也沒做。因此，幾十年來計算機科學都沒有注意到這個問題，我并不感到驚訝。到這里，AlphaDev 的工作原理也應該變得更加清晰了。從根本上說，DeepMind 構(gòu)建了一個人工智能，用于檢查匯編代碼，并隨機刪除一些代碼，看看代碼是否會出問題。我這么說并不是要否定 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 指令集，針對此類問題生成更精簡的代碼。可以看到，此處我們不需要任何創(chuàng)建臨時變量的指令：
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 硬件上更好地工作。我編寫了所有這類的匯編優(yōu)化，將性能提升到與在 x86 上運行 OpenSSL 相同的水平。MbedTLS 是簡單的、可移植的、可修改的 C 代碼，因此對于任何想要一個簡明易懂的匯編加密庫的人來說，這都是個好消息。我將自己的做法告訴了 Arm，雖然他們并不覺得有顛覆性，但仍然給予了友善的鼓勵。我希望有一天抽出時間學習 DeepMind 的做法，修改我的上游代碼。Arm 的 Optimized Routines 庫也非常豐富，該庫的雙精度數(shù)轉(zhuǎn)換在質(zhì)量上無可挑剔。對于 C 庫來說，這個庫的幫助特別大，因為幾十年來，開源社區(qū)一直依靠 Sun Microsystems 于 90 代初編寫的數(shù)學函數(shù)。Arm 找到了改進其中幾個函數(shù)的方法，例如 pow(x,y)。這可是最基本的數(shù)學運算之一，因此可以想象其影響力之大。例如，如果你使用 Arm 的解決方案在 x86 機器上實現(xiàn) pow(x,y)，那么執(zhí)行相同操作的速度將比英特爾的原生 x87 指令快 5 倍。很幸運 DeepMind 也加入了這個游戲，我冒昧地將他們的 libcxx diff 轉(zhuǎn)換成了簡單易讀的 C 代碼，這樣每個人都能欣賞到它的美麗。這是我希望 DeepMind 的論文和博客文章改進的另一個地方，因為你會在這段代碼中發(fā)現(xiàn)專家用來讓編譯器生成無分支 MOVcc 指令的規(guī)范技巧。
// 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 sortedstatic 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() 函數(shù)后，我覺得自己對 DeepMind 研究的動機有了更好的理解。在 ARM64 上編譯 Sort5() 函數(shù)，編譯器將生成一個包含 11 個寄存器的函數(shù)。你在推導一個數(shù)學方程式時，大腦里能否時同時思考 11 個變量？可能不行，這就是我們依賴 PartialSort3 這類內(nèi)核函數(shù)的原因。作為有感知力的生物，人類與猴子并沒有太大區(qū)別。我們變得更加聰明的主要因素是我們能夠解決難題，并將其分解為更小的問題。因此，很高興看到深度學習應用于增強我們的抽象能力。此外，還有一點值得一提，Sort3() 和 Sort5() 本身就是內(nèi)核，因為它們的目標就是成為傳統(tǒng)排序功能的構(gòu)建塊。DeepMind 的博客文章涵蓋了這個主題，但我認為分享一些可移植和可執(zhí)行的代碼可能會很有幫助。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 的新功能和改進。基本上就是快速排序，只是在遞歸到較小的切片時切換到排序內(nèi)核和插入排序。對于 libcxx，我認為他們甚至在堆排序中多加了一個步驟，雖然有點慢，但可以防止惡意代碼破壞棧。此時，可能你最想知道的是，我可以使用這個排序算法嗎？這些排序網(wǎng)絡內(nèi)核真的能提高排序速度嗎？我認為答案需要視情況而定。如果你只想對 long 進行升序排序，上面的代碼將比 C 庫提供的標準 qsort() 函數(shù)快 2 倍。而且你還不需要內(nèi)核來處理。到目前為止，我確定的是，在我的個人計算機（搭載了英特爾酷睿 i9-12900KS）上，上述函數(shù)對 long 進行排序的速度為每秒 255 兆字節(jié)。但是，如果我注釋掉排序內(nèi)核：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() 函數(shù)的排序速度可以達到每秒 275 兆字節(jié)。我在簡化算法后，又實現(xiàn)了 7% 的性能提升。這個函數(shù)就是 libc 在加載可執(zhí)行文件時對 elf 符號表進行排序時使用的函數(shù)。long 的好處是，它的長度足以存儲一個 int 鍵值對。能夠快速排序映射項是非常有用的技巧。結(jié)合樸素快速排序與樸素插入排序是我迄今為止找到的最佳解決方案，因為我必須平衡規(guī)模和性能。上面的函數(shù)編譯后只有 181 字節(jié)的 x86-64 機器碼。由于 DeepMind 的 sort3() 只有 42 字節(jié)，我希望我可以犧牲一些大小來獲得性能優(yōu)勢。因為到目前為止我發(fā)現(xiàn)的第二佳算法是基數(shù)排序，性能可達每秒 400 MB，除了依賴于 malloc() 之外，還需要高達 763 字節(jié)的二進制。所以很高興看到這些內(nèi)核的表現(xiàn)更好。我并不是說 DeepMind 的想法沒有價值。我認為值得注意的是，去年 DeepMind 非?？犊毓_了矢量化快速排序庫（當時還是 Google Brain），并以此實現(xiàn)了永遠無法挑戰(zhàn)的排序霸權(quán)。在我的電腦上，使用 Vqsort 對 long 進行排序的速度為每秒 1155 MB，甚至略勝于 djbsort，后者是開源社區(qū)中最受歡迎的庫之一，盡管除了 int 之外并沒有推廣至更多的數(shù)據(jù)類型。兩種實現(xiàn)方式都是通過向量化排序網(wǎng)絡來實現(xiàn)的。我認為這就是排序網(wǎng)絡技術大放異彩的地方。我想，如果不是 AlphaDev 的智能實體還處于起步階段，它也會采用這種做法。如果從第一原則開始，僅支持基礎指令集就非常困難。我認為再等一等，我們可以期待 AlphaDev 未來會取得偉大的成就，因為它會努力應對更艱巨的挑戰(zhàn)。此外，DeepMind 壓縮了算法的規(guī)模，這一點我也很喜歡，因為這并不常見。極限規(guī)模編程是我的愛好之一。我曾在一篇博客文章中發(fā)布過一個用于 lambda 演算的虛擬機只有 383 字節(jié)，還有一個擁有垃圾收集功能的 lisp 機器，只有 436 字節(jié)。我也曾在博客中介紹優(yōu)化 cosmpolitan c 庫大小的技巧。我也喜歡 DeepMind 的母公司谷歌，幾周前谷歌授予了我一份開源伙伴的獎金，很高興看到他們也對壓縮代碼規(guī)模充滿了熱情。很高興看到他們用我的庫來改進向量化快速排序。我只是希望全世界最佳 long 排序不要因為多加了 24 KB 的二進制編碼而變成 C++ 龐然大物。僅升序排序就有 23,000 行匯編代碼。我迫切希望有一天看到 AlphaDev 能夠?qū)@些代碼下手。最后，我喜歡一家人工智能公司建造用機器語言編寫代碼的機器的想法。為什么他們不喜歡這個想法呢？成為機器是機器的本性。作為一名開發(fā)人員，我發(fā)現(xiàn) OpenAI 正在創(chuàng)造的未來缺乏這樣的想法，他們建立了一個偉大的大型家長式機器，在零和經(jīng)濟中與地球上的每一位開發(fā)者競爭，然后吸引世界各地的租客通過政府監(jiān)管來控制那臺機器。OpenAI 承諾要自動化所有的任務，包括我最喜歡的編程，但他們正在努力創(chuàng)造的未來并不是在朝著這個方向前進。我想要的是能夠控制一臺能夠完成我自己無法完成的事情的機器，比如發(fā)現(xiàn)排序內(nèi)核。這才是真正的進步，我認為我們可以去除的每一行匯編代碼都是朝著這個夢想積極邁出的一步。
*博客內(nèi)容為網(wǎng)友個人發(fā)布，僅代表博主個人觀點，如有侵權(quán)請聯(lián)系工作人員刪除。
