AlphaDev將排序算法提速70%!C語(yǔ)言庫(kù)作者一文詳解DeepMind最新AI
幾天前,DeepMind推出了AlphaDev,直接把排序算法提速70%。
這一全新AI系統(tǒng),便是基于下棋高手AlphaGo打造。
而這項(xiàng)研究恰恰激起了前谷歌研究人員Justine Tunney的興趣。
她表示,作為一名C語(yǔ)言庫(kù)的作者,我一直在尋找機(jī)會(huì)來(lái)策劃最好的東西。
一起看看Justine如何詳解DeepMind排序算法。
DeepMind排序算法
DeepMind的這一發(fā)現(xiàn)贏得了當(dāng)之無(wú)愧的關(guān)注,但不幸的是,他們本可以更好地解釋AlphaDev。
接下來(lái),從DeepMind發(fā)布的匯編代碼開(kāi)始,該代碼將一個(gè)有三個(gè)項(xiàng)目的數(shù)組進(jìn)行排序,從偽匯編翻譯成匯編:
我將這個(gè)函數(shù)命名為 move37() ,是因?yàn)镈eepMind的博客文章,將其與AlphaGo下的令人震驚的「第37步」進(jìn)行了比較。
在2016那場(chǎng)人機(jī)大戰(zhàn)中,AlphaGo下了一顆違反人類直覺(jué)的棋,一個(gè)簡(jiǎn)單的肩沖,擊敗了傳奇圍棋選手李世石。
所以如果運(yùn)行DeepMind代碼:
但是,在我看來(lái)這是一個(gè)錯(cuò)誤。
我們給它的數(shù)組是{3,1,2},但 move37() 將其排序?yàn)閧2,1,3}。
DeepMind一定在欺騙我們,因?yàn)槲也幌嘈?在1之前。再來(lái)看看他們對(duì)LLVM libcxx所做的開(kāi)源貢獻(xiàn),這有望澄清一些事情:
所以 move37() 實(shí)際上不是一個(gè)排序函數(shù),而是一個(gè)排序內(nèi)核,旨在用作 sort3() 函數(shù)的構(gòu)建塊。
如果論文和博客文章能提到這一點(diǎn)就好了,因?yàn)樗屛以谧疃痰臅r(shí)間內(nèi)感到非常困惑。下面是更好的代碼版本,其中包括缺失的交換(swap)操作。
為了解釋為什么他們的代碼很重要,讓我們考慮一下這個(gè)算法在高層次上是如何工作的。當(dāng)我第一次嘗試自己解決 sort3() 問(wèn)題時(shí),我想到了這個(gè):
然后我查看了libcxx,發(fā)現(xiàn)它們也在做同樣的事情。上述代碼的問(wèn)題是,編譯器并不善于優(yōu)化它。
如果你嘗試編譯上面的代碼,就會(huì)注意到你的編譯器插入了大量的分支指令。這就是DeepMind試圖通過(guò)LLVM貢獻(xiàn)來(lái)改進(jìn)的地方。
然而,這些技術(shù)往往不太容易理解。
我實(shí)際上喜歡天真無(wú)邪的代碼,因?yàn)槿绻覀儾[起眼睛,可以看到一種模式,與DeepMind最先進(jìn)的匯編代碼有相同的基本想法。
這個(gè)想法是這個(gè)問(wèn)題本質(zhì)上歸結(jié)為3個(gè)比較和交換操作:
上面的代碼是之前排序網(wǎng)絡(luò)的最先進(jìn)技術(shù)?,F(xiàn)在,這就是DeepMind的新發(fā)現(xiàn)發(fā)揮作用的地方。他們發(fā)現(xiàn)有時(shí)上面的 mov 指令是不必要的。
如果你試著運(yùn)行上面的代碼,你會(huì)發(fā)現(xiàn)不管有沒(méi)有被刪除的行,它都是100%正確的。
這行代碼看起來(lái)像是在做什么,但實(shí)際上什么也沒(méi)做。所以我并不驚訝這樣的事情會(huì)被計(jì)算機(jī)科學(xué)忽視幾十年。
現(xiàn)在也應(yīng)該更清楚AlphaDev是如何工作的。
DeepMind基本上構(gòu)建了一個(gè)人工智能,它可以擺弄匯編代碼,隨機(jī)刪除一些東西,看看它是否損壞。
我這么說(shuō)并不是要否定AlphaDev的智能,因?yàn)槿绻艺f(shuō)我沒(méi)有做同樣的事情,那就是在撒謊。
上面的代碼中還有兩個(gè) mov 指令,我們有可能將其刪除。通過(guò)使用ARM64指令集來(lái)做到這一點(diǎn),它可以為類似的問(wèn)題提供更小的代碼。
在這里,我們不需要任何指令來(lái)創(chuàng)建臨時(shí)變量:
Arm公司最近風(fēng)頭正勁,我想上面的例子可以作為他們贏得名聲的證據(jù)。
Arm也是目前開(kāi)源領(lǐng)域最好的公司之一。比如,他們的MbedTLS庫(kù)是我迄今為止見(jiàn)過(guò)的最被低估的瑰寶之一。
當(dāng)我開(kāi)始使用它時(shí),我原本有這樣的計(jì)劃,即修改Arm的代碼,使之在x86硬件上更好地工作。
我編寫了所有這些精心設(shè)計(jì)的匯編優(yōu)化,使其與x86上的OpenSSL達(dá)到相同的性能。
MbedTLS是簡(jiǎn)單、可移植、可破除的C代碼,因此對(duì)于任何想要一個(gè)不是Perl生成的匯編的加密庫(kù)的人來(lái)說(shuō),是個(gè)好消息。
我告訴了Arm公司的人我在做什么,他們并沒(méi)有覺(jué)得這是顛覆性的。
我希望有一天能找到時(shí)間做DeepMind做的事情,并在上游進(jìn)行修改。Arm公司的優(yōu)化程序庫(kù)也是多產(chǎn)的,它在質(zhì)量上與雙轉(zhuǎn)換無(wú)懈可擊。
它對(duì)C庫(kù)對(duì)此特別感興趣,因?yàn)閹资陙?lái),開(kāi)源社區(qū)一直依靠Sun Microsystems在90年代初編寫的數(shù)學(xué)函數(shù)來(lái)維持生計(jì)。
Arm找到了一種改進(jìn)其中幾個(gè)函數(shù)的方法,例如 pow(x,y) ??紤]到這是數(shù)學(xué)中最基本的運(yùn)算之一,這是一件非常有影響力的事情。
比如,如果你在純軟件中使用Arm的解決方案在x86機(jī)器上實(shí)現(xiàn) pow(x,y) ,那么它將比英特爾的原生x87指令快5倍。
很幸運(yùn),DeepMind也加入了這個(gè)游戲,所以我冒昧地把他們的libcxx diff翻譯成可讀的C代碼。
這是我希望在論文和博客文章中看到的另一件事,因?yàn)樵谶@段代碼中,你會(huì)發(fā)現(xiàn)專家們用來(lái)讓編譯器生成無(wú)分支 MOVcc 指令的規(guī)范技巧。
當(dāng)我看到 Sort5() 函數(shù),我覺(jué)得自己對(duì)DeepMind研究的動(dòng)機(jī)有了更好的理解。
如果你在ARM64上編譯 Sort5() 函數(shù),那么編譯器將產(chǎn)生一個(gè)處理11個(gè)寄存器的函數(shù)。如果你在推理一個(gè)數(shù)學(xué)方程,那么你能一次在你的工作記憶中保存11個(gè)變量嗎?
可能不會(huì)。這就是為什么有一個(gè)像 PartialSort3 這樣優(yōu)秀的內(nèi)核函數(shù)如此有用的原因。
值得一提的是, Sort3() 和 Sort5() 本身就是內(nèi)核,因?yàn)樗鼈冎荚诔蔀閭鹘y(tǒng)排序功能的構(gòu)建塊。
博客文章涵蓋了這個(gè)主題,但我認(rèn)為分享一些實(shí)際上可移植和可執(zhí)行的東西會(huì)很有用。
The above algorithm shows what the new and improved libcxx is doing. It's basically quicksort except it switches to the sorting kernels and insertion sort when recursing into smaller slices. With libcxx I think they even took the added step of schlepping in heapsort, which is kind of slow, but prevents adversaries from smashing your stack. 上面的算法顯示了新的和改進(jìn)的libcxx正在做什么。它基本上是快速排序,除了在遞歸到更小的切片時(shí)切換到排序內(nèi)核和插入排序。對(duì)于libcxx,我認(rèn)為他們甚至采取了在堆排序中移動(dòng)的額外步驟,這有點(diǎn)慢,但可以防止對(duì)手破壞您的堆棧。
The main thing you may be wondering at this point is, can I use this? Do these sorting network kernels actually make sorting go faster? I would say yes and no. When all you want is to sort ascending longs, the code above will go 2x faster than the standard qsort() function provided by your C library. Except you don't need the kernels to do that. What I've determined so far is that, on my personal computer (which has an Intel Core i9-12900KS) the above function sorts longs at 255 megabytes per second. However if I comment out the sorting kernels: 在這一點(diǎn)上,你可能想知道的主要事情是,我可以使用這個(gè)嗎?這些排序網(wǎng)絡(luò)內(nèi)核真的能讓排序變得更快嗎?我會(huì)說(shuō)是和不是。當(dāng)你只想對(duì)升序長(zhǎng)進(jìn)行排序時(shí),上面的代碼將比你的C庫(kù)提供的標(biāo)準(zhǔn) qsort() 函數(shù)快2倍。只是你不需要內(nèi)核來(lái)做到這一點(diǎn)。到目前為止,我已經(jīng)確定,在我的個(gè)人電腦上(它有一個(gè)英特爾酷睿i9-12900KS),上面的函數(shù)以每秒255兆字節(jié)的速度排序。但是如果我注釋掉排序內(nèi)核:
然后我的 longsort() 函數(shù)以每秒275兆字節(jié)的速度運(yùn)行,通過(guò)簡(jiǎn)化算法實(shí)現(xiàn)了7%的性能提升。
long 的好處是它足夠長(zhǎng),可以存儲(chǔ) int 鍵值對(duì),能夠快速對(duì)地圖條目進(jìn)行排序是一個(gè)有用的技巧。
上面的函數(shù)編譯后只有181字節(jié)的x86-64機(jī)器代碼。
由于DeepMind的 sort3() 只有42字節(jié),我希望可以交換一些大小以獲得性能優(yōu)勢(shì)。
因?yàn)榈侥壳盀橹?,我發(fā)現(xiàn)的下一個(gè)最佳算法是改用基數(shù)排序,速度為400 MB/s,但除了依賴于 malloc() 之外,還需要高達(dá)763字節(jié)的二進(jìn)制占用空間。因此,如果能看到這些內(nèi)核做得更好就好了。
這并不是說(shuō)DeepMind的想法沒(méi)有價(jià)值。
我認(rèn)為值得注意的是,DeepMind非??犊?,去年給了我們他們的矢量化快速排序庫(kù)(當(dāng)時(shí)他們被稱為Google Brain),并通過(guò)這樣做實(shí)現(xiàn)了永遠(yuǎn)無(wú)法挑戰(zhàn)的排序優(yōu)勢(shì)。
Vqsor在我的電腦上以1155 MB/s的速度對(duì)長(zhǎng)時(shí)間進(jìn)行排序。
它甚至略微優(yōu)于djbsor,后者是開(kāi)源社區(qū)中最受歡迎的庫(kù)之一,盡管它從未推廣到比 int 更多的數(shù)據(jù)類型。
這兩種實(shí)現(xiàn)實(shí)現(xiàn)的方式都是通過(guò)矢量化排序網(wǎng)絡(luò)。我認(rèn)為這就是排序網(wǎng)絡(luò)技術(shù)真正閃耀的地方。
我想,如果就智能實(shí)體而言,AlphaDev不是一個(gè)蹣跚學(xué)步的孩子,它就會(huì)這樣做。
當(dāng)你從基本原則開(kāi)始時(shí),僅基線指令集就非常難以支持。如果我們等待,那么我認(rèn)為我們可以期待在未來(lái)看到AlphaDev的偉大成就,因?yàn)樗谂?yīng)對(duì)更強(qiáng)大的挑戰(zhàn)。
我也很喜歡DeepMind讓算法變得更小的事實(shí),因?yàn)檫@是我不??吹降摹?/span>
大小編碼是我最喜歡的愛(ài)好之一。在這個(gè)博客上,我發(fā)布了一個(gè)383字節(jié)的lambda演算虛擬機(jī)和一個(gè)436字節(jié)的帶有垃圾回收機(jī)制的lisp機(jī)。
我還在博客上介紹了我在cosmpolitan c庫(kù)中使用的大小優(yōu)化技巧。
我也喜歡DeepMind的母公司,因?yàn)閹字芮癎oogle給我頒發(fā)了開(kāi)源同行獎(jiǎng)金,很高興看到他們分享我使軟件變小的熱情。
很高興看到他們用它來(lái)改進(jìn)矢量化快速排序。
最后,我喜歡人工智能公司用機(jī)器語(yǔ)言編寫代碼的機(jī)器的想法。他們?yōu)槭裁床荒??機(jī)器的本質(zhì)就是機(jī)器。
作為一個(gè)建設(shè)者,我發(fā)現(xiàn)這比OpenAI正在創(chuàng)造的未來(lái)要少得多。
他們已經(jīng)建立了一個(gè)巨大的家長(zhǎng)式機(jī)器,在零和經(jīng)濟(jì)中與地球上的每個(gè)建設(shè)者競(jìng)爭(zhēng),然后誘使世界上的尋租者通過(guò)政府監(jiān)管來(lái)控制這臺(tái)機(jī)器。
我不認(rèn)為OpenAI承諾將所有我最喜歡做的任務(wù)(如編碼)自動(dòng)化是一種進(jìn)步。我想要的是能夠控制一臺(tái)機(jī)器,這臺(tái)機(jī)器能夠完成我自己無(wú)法完成的事情,比如發(fā)現(xiàn)排序內(nèi)核。這才是真正的進(jìn)步。
我認(rèn)為,我們能夠砍掉的每一條裝配線都是朝著這個(gè)夢(mèng)想的積極方向邁出的一步。