前言

DeepMind 最近在 Nature 發(fā)表了一篇論文 AlphaDev[2, 3]，一個利用強化學習來探索更優(yōu)排序算法的AI系統(tǒng)。

AlphaDev 系統(tǒng)直接從 CPU 匯編指令的層面入手去探索更優(yōu)的排序算法，因為相對于高級編程語言來說，在匯編指令層級對存儲和寄存器的操作可以更加的靈活，所以能發(fā)現(xiàn)更多潛在的調(diào)優(yōu)策略。

在 AlphaDev 的論文中，只關注探索短序列排序：

• 定長序列排序（比如 sort3 算法只能對長度為3的序列進行排序）
• 變長序列排序（比如 variable sort5 算法可以對長度為1~5的變長序列進行排序）

而對于長序列的排序，可以被分解為短序列的排序。

DeepMind 通過 AlphaDev 發(fā)現(xiàn)了比目前人工調(diào)優(yōu)算法更優(yōu)的定長短序列排序算法 sort3，sort4 和 sort5 ，并且已經(jīng)將代碼提交到了 LLVM 標準 C++ 庫[4] 。

簡單來說，AlphaDev 將探索更高效排序算法的過程，建模為一個單玩家的匯編游戲（single-player game， AssemblyGame）。

游戲的過程就是玩家從 CPU 匯編指令集合中，選取一系列的指令組合得到一個新的排序算法。不過這個過程是非常有挑戰(zhàn)的，玩家需要考慮，匯編指令的組合空間并最終得得到一個正確和高效的算法。

該游戲主要包括以下難點：

• 匯編游戲的搜索空間和圍棋類似（10^700）
• 只要有一條指令沒弄對，可能就會導致整個算法錯誤

AlphaDev 系統(tǒng)詳解

將排序算法表示為 CPU 匯編指令

首先來看一個簡單的變長（variable sort2）短排序函數(shù)的 C 代碼實現(xiàn)，排序結(jié)果從小到大：

voidvariable_sort_2(intlength,int*a){
switch(length){
case0:
case1:
return;
case2:
inttmp=a[0];
//a[0]保存兩者之間的最小值
a[0]=(a[1]0])?a[1]:a[0];
//a[1]保存兩者之間的最大值
a[1]=(a[1]1];
return;
}
}

通過 gcc 生成對應的匯編代碼，我用的 gcc 版本是 11.3.0，命令 gcc -S -O1 -o sort2.s sort2.c

匯編代碼只保留了核心部分，生成的結(jié)果和論文中的示例有些許不同但是原理是一致的：

variable_sort_2:  
.LFB0:
; %edi 寄存器保存參數(shù) length 的值
; cmpl 指令對比 %edi 和 常量 2
cmpl$2, %edi 
; 相等就跳轉(zhuǎn)到 .L3 標簽處，
        ; 對應 C 代碼的 case 2
je.L3
.L1:
; 不等于 2 就直接返回，
        ; 對應 C 代碼 case 0 和 1
ret 
.L3:
; 將 a[0] 賦值給寄存器 %edx 
movl(%rsi), %edx
; 將 a[1] 賦值給寄存器 %eax 
movl4(%rsi), %eax
; 對比 %edx 和 %eax
cmpl%edx, %eax
; 將 %edx 賦值給 %ecx
movl%edx, %ecx
; cmov 是條件移動指令根據(jù) cmpl 
; 指令的結(jié)果判斷是否執(zhí)行
; 如果 %eax <= %edx 
; 則將 %eax 賦值給 %ecx
cmovle%eax, %ecx
; 此時 %ecx 保存了最小值
; 將 %ecx 賦值給 a[0]
movl%ecx, (%rsi)
; 如果 %eax 小于 %edx
; 則將 %edx 賦值給 %eax
cmovl%edx, %eax
; 此時 %eax 保存了最大值
; 將 %eax 賦值給 a[1]
movl%eax, 4(%rsi)
jmp.L1

一般來說匯編程序所做的事情基本都是，將內(nèi)存的值復制到寄存器，然后對寄存器的值作修改，再將寄存器的值寫回到內(nèi)存中。

而 AlphaDev 系統(tǒng)只關注 x86 處理器架構所支持的匯編指令集合的一個子集。

每條匯編指令的格式均為：操作碼<操作數(shù)A, 操作數(shù)B> 比如：

• mov 移動指令，表示將 A 的值賦值給 B

• cmp 比較指令，相當于 執(zhí)行 A - B 操作，但是不會對 A 和 B 做修改，而是根據(jù)相減的結(jié)果設置特殊的 flag 寄存器，更多內(nèi)容可以參考[5]

• cmovX 條件移動指令，根據(jù) X 和 flag 寄存器的值判斷是否執(zhí)行將 A 賦值給 B 的操作，一般都是出現(xiàn)在 cmp 指令之后。X 可以是 L （是否滿足小于條件）， G （是否滿足大于條件），LE （是否滿足小于或等于條件），GE （是否滿足大于等于條件）。

• jX 條件跳轉(zhuǎn)指令，根據(jù) X 和 flag 寄存器的值判斷是否執(zhí)行跳轉(zhuǎn)到指定標記位置操作，A 可以是匯編程序代碼中的標記位置，如上面所示匯編代碼的 .L1 和 .L3。X 可以是 NE （是否不等于），E （是否等于）或者可以填表示無條件跳轉(zhuǎn)。

將探索更優(yōu)排序算法表示為強化學習問題

AlphaDev 將 CPU 匯編指令層面的算法優(yōu)化過程轉(zhuǎn)化為一個單玩家的游戲。

游戲每一步的狀態(tài)定義為 : St = 。

其中， Pt 表示游戲到至今為止所生成的算法，Zt 則表示在給定輸入的前提下執(zhí)行完 Pt 里的指令之后，內(nèi)存和寄存器的狀態(tài)。

如上圖所示，在時間步 t ，AlphaDev 接受到當前狀態(tài) St 和 所要執(zhí)行的動作 at （比如 mov ），也就是往當前生成的算法 Pt 中添加的合法匯編指令。

在添加完指令之后，就是計算獎勵分數(shù) rt （包括評估算法的正確性和延遲）。

算法正確性評估

正確性評估就是將 N 組測試序列輸入到算法 Pt 中，得到N 組輸出，和正確的排序結(jié)果最比較來計算獎勵分數(shù)。

論文中給出了3種正確性評估函數(shù)，首先定義 P 為輸入序列長度， PCt 為在時間步 t 序列中，位置正確的值的個數(shù)，這里我理解應該是和正確的排序結(jié)果逐個位置對比，統(tǒng)計相等的個數(shù)。

三個函數(shù)分別定義如下：

• func1 = (P - PCt) / P
• func2 = sqrt(func1)
• func3 = sqrt(PCt)

論文中提到采用第三個函數(shù)效果最好。

延遲評估

延遲分數(shù)的計算可以是：

• 對系統(tǒng)增加代碼長度計算懲罰，因為代碼的長度一般都是和耗時高度相關
• 直接計算算法的真實耗時

整個強化學習的游戲在執(zhí)行有限步驟之后就會被終止。只有生成正確而又低延遲的匯編代碼才算贏得游戲。而不管是生成了錯誤的代碼還是正確但低效的實現(xiàn)都視為游戲輸了。

AlphaDev 采用的強化學習算法是對 AlphqaZero 算法的擴展，也是采用深度神經(jīng)網(wǎng)絡來引導蒙特卡洛樹搜索（MCTS）的規(guī)劃過程。網(wǎng)絡模型的輸入是 St ，輸出是對動作策略和獎勵的預測。

整個游戲過程簡單來說就是，用一個固定參數(shù)的網(wǎng)絡模型，通過給定的當前狀態(tài)執(zhí)行一個蒙特卡洛樹搜索過程，然后采取下一步動作。然后可以用生成的游戲過程（包含每一步的狀態(tài)和獎勵）去訓練和更新網(wǎng)絡的參數(shù)。

網(wǎng)絡模型結(jié)構

模型包含兩部分：

• 一個 Transformer 編碼器模塊，用于建模算法，輸入是至今為止生成的匯編指令序列
• 一個 CPU 狀態(tài)編碼器 MLP 模塊，輸入當前寄存器和內(nèi)存的狀態(tài)

兩個網(wǎng)絡的輸出 embedding 會合并在一起來表示當前的狀態(tài)。

網(wǎng)絡模型整體的結(jié)構如下：

Transformer 編碼器模塊具體圖示

如上圖所示，把當前生成的匯編代碼序列的每一條指令的操作碼和操作數(shù)都轉(zhuǎn)換為 one-hot 編碼序列，然后輸入到網(wǎng)絡中。

但是具體的 one-hot 編碼規(guī)則、詞表怎么設置、還有對于 CPU 狀態(tài)編碼網(wǎng)絡寄存器和內(nèi)存的狀態(tài)是怎么表示為網(wǎng)絡的輸入的等等，這些細節(jié)我在論文里沒找到。

然后兩個網(wǎng)絡的輸出 embedding 會合并到一起接著輸入到幾個函數(shù)頭里計算，分別是預測下一步策略的函數(shù)頭，預測算法正確性的函數(shù)頭和預測算法真實延遲的函數(shù)頭。

網(wǎng)絡參數(shù)超參設置

論文的補充資料中提供了網(wǎng)絡的參數(shù)和三個函數(shù)頭的具體配置。

而對于策略的預測，論文中提到為了簡化問題和提高收斂性，而對動作空間做了一些限制，規(guī)則如下：

• 必須按照升序方式讀取內(nèi)存
• 寄存器按照升序分配
• cmp 和 cmovX 指令的操作數(shù)不能出現(xiàn)內(nèi)存地址
• 對每個內(nèi)存位置，只能讀取和寫入一次
• 每個寄存器在使用之前，必須初始化
• 不能連續(xù)調(diào)用 cmp 指令

訓練細節(jié)

AlphaDev 的訓練采用了 TPU v3，每個 TPU 核的 batch size 是 1024 ，總共用了 16 個 TPU 核，總共訓練了 100 萬次迭代。而在對于玩游戲積累訓練數(shù)據(jù)來說，則是在 TPU v4 上進行，總共用了 512 個 TPU 核。

實驗結(jié)果表明，最多只需2天模型就能訓收斂。

實驗結(jié)果

生成的算法和人工調(diào)優(yōu)對比

從實驗結(jié)果表格可以看到，對于短序列排序算法 AlphaDev 生成的代碼長度更短，而且平均耗時也更低。

對生成算法延遲的評估方式，比如對于 sort3 則是在 100 臺機器上做評估，每臺機器隨機生成 1000 條 3個數(shù)的序列，然后每條序列輸入到算法中，對這 1000 次評估取第5百分位數(shù)作為最終的評估結(jié)果（排除 cache miss 和 任務搶占 等因素）。

耗時采用的是 CPU_CLK_UNHALTED.CORE 這個計數(shù)器結(jié)果， 其計數(shù)值表示在一個特定時間段內(nèi)，處理器內(nèi)核的時鐘周期數(shù)。這個值越高，意味著處理器內(nèi)核在該時間段內(nèi)執(zhí)行了更多的指令。

AlphaDev 發(fā)現(xiàn)新的算法

對于定長序列排序，當應用到排序網(wǎng)絡算法[6]（sorting network algorithm）的時候 AlphaDev 生成的代碼中包含了一些有趣指令序列，相對于原始指令序列可以減少一條匯編指令，論文中稱之為：

• AlphaDev swap move
• AlphaDev copy move

啥是排序網(wǎng)絡算法？

排序網(wǎng)絡算法（Sorting Network Algorithm）是一種能夠?qū)σ唤M輸入數(shù)據(jù)進行排序的并行算法，其具有較好的并行性能適用于多處理器或多核心系統(tǒng)。

該算法的特點是，它將所有的比較和交換操作預先規(guī)劃好形成一個固定的結(jié)構，然后將輸入數(shù)據(jù)按照這個結(jié)構進行排序。

排序網(wǎng)絡由比較器（comparator）和線（wire）組成，如下圖所示：

水平線表示 wire，每條水平線持有一個待排序的值。兩條 wire 之間的垂直線段就表示一個比較器，比較器對比兩條水平線的值，如果比較器下方的值小于上方的值則交換兩條橫線的值，否則則不交換。

一個優(yōu)化過的排序網(wǎng)絡可以以最少的比較器，并將這些比較器放置在特定位置上，來實現(xiàn)對任意序列進行排序。

下圖是對一個構造好的排序網(wǎng)絡，輸入真實待排序序列的例子：

可見初始輸入是 [2, 3, 1, 4]，這些隨機數(shù)從左到右按順序經(jīng)過這些比較器之后，就得到了排序好的序列 [1, 2, 3, 4]。

AlphaDev swap move

先來看這個排序網(wǎng)絡，只看紅圈部分的功能就是對給定的輸入 [A, B, C] 將其轉(zhuǎn)換為 [min(A,B,C), max(min(A,C),B), max(A,C)]。

然后經(jīng)過 AlphaDev 優(yōu)化之后，可以將第一個輸出的 min(A,B,C) 改為只計算 min(A,B)，原因是因為前面的 B 和 C橫線之間經(jīng)過比較器之后已經(jīng)有了前置條件 B <= C。

而通過這個優(yōu)化就能省去一條匯編指令，下圖是紅圈部分的偽代碼實現(xiàn)：

左邊是原始偽代碼實現(xiàn)，右邊是經(jīng)過 AlphaDev 優(yōu)化之后的實現(xiàn)，可以看到少了一條匯編指令 mov S P。

AlphaDev copy move

接下來看對4個元素進行排序的排序網(wǎng)絡，是在對 sort8 這個算法優(yōu)化過程中發(fā)現(xiàn)的。該排序網(wǎng)絡對于輸入序列 [A, B, C, D] 轉(zhuǎn)換為 [min(A, B, C, D), max(B, min(A, C, D), max(C, min(A, D)), max(A, D) ]。

該排序網(wǎng)絡是 sort8 的一個子排序網(wǎng)絡，而根據(jù)比較器的放置位置來看，A 和 D 比較之后后續(xù)就不再和其他元素比較了，所以D出來的結(jié)果就是四個元素中最大的，所以隱含了一個條件就是 D >= min(A, C)。

因此對第二個輸出元素的計算可以從 max(B, min(A, C, D)) 改為 max(B, min(A, C))，就可以節(jié)省一條匯編指令。

偽代碼如下:

左邊是原始偽代碼實現(xiàn)，右邊是經(jīng)過 AlphaDev 優(yōu)化之后的實現(xiàn)，可以看到少了一條匯編指令 mov P T。

總結(jié)

這篇文章只是對 AlphaDev 論文中的主要內(nèi)容作解讀，對于更多的內(nèi)容和細節(jié)感興趣的讀者可以查閱原論文和論文的補充資料 [2,3]，DeepMind 也也開源了一份偽代碼實現(xiàn) [7]。