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