基于強化學(xué)習(xí)的循環(huán)自動展開研究

李居垚，何先波

（西華師范大學(xué)電子信息工程學(xué)院，四川 南充 637009）

0 引 言

當(dāng)下新架構(gòu)芯片問世的頻率遠超以往，這需要縮短芯片編譯的開發(fā)周期來進行適配。 在編譯器中許多關(guān)鍵決策是由抽象模型來決定的，雖然這些模型能夠有效地解決問題，但是開發(fā)編寫一個精確的抽象模型需要耗費大量的人力物力；所以使用新技術(shù)來縮短編譯器開發(fā)的成本和周期是有必要的。

編譯優(yōu)化中的循環(huán)展開有著承上啟下的作用，因為循環(huán)展開從多維度間接影響了系統(tǒng)性能，同時也會為其他優(yōu)化（如SLP）帶來契機，所以開發(fā)一種有別于傳統(tǒng)的新技術(shù)來對編譯器循環(huán)展開進行自動化是必要的。 本文將使用人工智能技術(shù)提出一個模型來對編譯器循環(huán)展開工作中的至關(guān)重要的展開因子（Unroll Factor，UF）進行自動生成，已達到循環(huán)展開工作自動化減少編譯器開發(fā)工作量和成本的目的。

1 相關(guān)理論及方法

1.1 循環(huán)展開

循環(huán)展開是一種常用的編譯優(yōu)化策略，開啟后編譯器會通過既定的抽象模型啟發(fā)式地將目標(biāo)循環(huán)進行多次復(fù)制。 循環(huán)展開主要用于為其他優(yōu)化提供機會，例如，展開會創(chuàng)建與動態(tài)執(zhí)行相對應(yīng)的多個靜態(tài)內(nèi)存單一操作指令，可以重新安排這些指令以利用內(nèi)存局部性［1］。 如果循環(huán)訪問連續(xù)迭代中的相同內(nèi)存位置，引用還可以用標(biāo)量進行消除。 循環(huán)展開是暴露相鄰內(nèi)存引用的關(guān)鍵，展開完成后后續(xù)的優(yōu)化可識別到這些引用并將其合并為一個寬引用［2］。

由于循環(huán)展開的主要影響集中在次級效應(yīng)上，所以選取最優(yōu)的展開因子進行循環(huán)展開是困難的。循環(huán)展開也可以被歸納為空間換時間的一種手段，這很容易造成一種循環(huán)展開總能提升性能的誤解。不合適的循環(huán)展開都會造成以下性能損失：會降低指令緩存的性能，會造成內(nèi)存的溢出和重新加載，如果編譯器選擇推測性的訪問展開內(nèi)存會造成動態(tài)沖突［3］。 循環(huán)展開會增加編譯時間，所以通過暴力搜索等手段獲取循環(huán)展開因子是低效的，圖1 展示了不同循環(huán)展開因子選擇對編譯耗時的影響。

圖1 TSVC 中s3111.c 設(shè)置不同循環(huán)展開次數(shù)的編譯耗時比Fig. 1 Compilation time ratio of s3111. c with different loop unwinding times in TSVC

使用暴力搜索的方式來尋找最優(yōu)的展開因子是效率低下的，如圖1 所示。 圖1 中，以不進行展開的耗時為基準(zhǔn)。 近年來，隨著機器學(xué)習(xí)技術(shù)不斷發(fā)展成熟，許多學(xué)者希望利用機器學(xué)習(xí)來解決復(fù)雜的編譯器優(yōu)化問題，即構(gòu)建一個機器學(xué)習(xí)模型去預(yù)測編譯器優(yōu)化參數(shù)和編譯選項等［4］。 本文將會嘗試使用強化學(xué)習(xí)來預(yù)測的循環(huán)展開因子。

1.2 強化學(xué)習(xí)

強化學(xué)習(xí)通過和環(huán)境動態(tài)交互不斷試錯來得到較優(yōu)解的方法。 在這種方法中智能體通過不斷和環(huán)境交互來進行學(xué)習(xí)［5］，強化學(xué)習(xí)想要直接訓(xùn)練得到策略函數(shù)或者價值函數(shù)是比較困難的，對此引入神經(jīng)網(wǎng)絡(luò)來近似目標(biāo)函數(shù)是一種常用的策略。

在強化學(xué)習(xí)中智能體通過對環(huán)境的“觀察”來采取行為動作，行為會產(chǎn)生新的環(huán)境并根據(jù)當(dāng)前環(huán)境得到一個獎勵的分值作為行為好壞的評判，整個學(xué)習(xí)過程的目的是獎勵達到數(shù)學(xué)期望的最大化并得到相映的決策函數(shù)（π?）：

強化學(xué)習(xí)有很多巧妙的算法，本文選用PPO（Proximal Policy Optimization）算法［6］。 PPO 算法的優(yōu)勢在于每次進行梯度更新時能夠使得成本最小化，并且能夠很好地消除和上一次梯度更新的偏差。

1.3 代碼嵌入

使用代碼詞嵌入能夠使輸入的源程序更好地通過強化學(xué)習(xí)進行訓(xùn)練，代碼詞嵌入的最終目標(biāo)是有效地將源程序和強化學(xué)習(xí)所需要的向量進行映射。在本文中，選擇IR2Vec 作為代碼嵌入。

IR2Vec［4］是一種基于源碼中間表示形式（LLVM IR）的代碼嵌入，這種代碼嵌入的方法具有簡明和可伸縮的特點，IR2Vec 通過表征學(xué)習(xí)和流信息相結(jié)合的方式捕獲程序的語法和語義將程序表示為連續(xù)空間中的分布式嵌入。 IR 的表示被轉(zhuǎn)化為一個種子嵌入詞匯表，并且有2 種編碼模式：符號編碼和流感知編碼。 由于本文主要針對的是源代碼中循環(huán)部分的工作，控制流的信息至關(guān)重要，所以本文使用IR2Vec 的流感知編碼模式。

2 相關(guān)工作

2.1 實驗平臺

實驗平臺信息見表1。

表1 實驗平臺信息Tab. 1 Experimental platform information

2.2 實驗內(nèi)容

本次實驗所使用的模型如圖2 所示，源程序作為初始的輸入會通過腳本識別代碼中的循環(huán)并插入循環(huán)展開指令設(shè)置循環(huán)展開因子，然后使用Clang＼LLVM 將源代碼轉(zhuǎn)為中間表示形式（IR），而后將IR輸入到IR2Vec 中生產(chǎn)長度為300 的向量嵌入到強化學(xué)習(xí)的智能體中，同時編譯運行使用LLVM 基準(zhǔn)（LLVM－O3）的可執(zhí)行文件并收集運行時間。 下一階段使用策略函數(shù)輸出的新因子作為參數(shù)，再次進行編譯并運行，2 次運行的性能差異將作為模型的獎勵，這意味著本模型的獎勵不需要手動進行設(shè)置，獎勵的值也不是固定不變的。 模型中一輪迭代后會將預(yù)測出的循環(huán)展開因子作為下一輪的學(xué)習(xí)所需要的編譯選項參數(shù)，獲得新的IR、新的向量嵌入智能體，以此往復(fù)訓(xùn)練出預(yù)測循環(huán)展開因子的模型。

圖2 本文提出的自動循環(huán)展開框架Fig. 2 The automatic loop unrolling framework proposed in this article

本次實驗將會采取上述的模型和方法進行，代碼嵌入使用的IR2Vec 開源工具并根據(jù)設(shè)定的任務(wù)要求進行改造，研究中選用其流感知模式來更好地識別循環(huán)中的控制流信息。 在搭建強化學(xué)習(xí)模型的過程中，使用了RLlib［6］和Tune［7］，RLlib 和Tune 都能在強化學(xué)習(xí)的開源庫里找到。 這樣一來，研究所搭建的模型會更加規(guī)范，便于進行參數(shù)的調(diào)整和對模型的擴展。

在實驗中使用Ray［8］框架，因此對于提前設(shè)置好的3 種不同的學(xué)習(xí)率可以同時進行，這極大地縮短了訓(xùn)練所需的時間。 同時，為了與本次研究搭建的強化學(xué)習(xí)模型進行對比，還準(zhǔn)備了暴力搜索來鎖定最優(yōu)循環(huán)展開因子。

2.3 參數(shù)設(shè)置

實驗選用PPO 算法，設(shè)置了3 個學(xué)習(xí)率：5－e3、5－e4、5－e5。 本模型的獎勵將以LLVM11 －O3 編譯后的運行時間作為基準(zhǔn)，將其進行如下定義：

對于循環(huán)展開因子的取值，本文設(shè)立一個數(shù)組array［1、2、4、…、8］，此后隨機從中取值作為強化學(xué)習(xí)的動作。

表2 是本文構(gòu)建的強化學(xué)習(xí)模型具體的參數(shù)設(shè)置。

表2 強化學(xué)習(xí)模型參數(shù)Tab. 2 Parameters of reinforcement learning model

2.4 數(shù)據(jù)集

實驗所使用的訓(xùn)練集是在向量化測試套件（TSVC）的基礎(chǔ)上進行篩選和拓展組成的訓(xùn)練集。TSVC 最初由Callahan 等人開發(fā)后由Maleki 等學(xué)者［9］將其從Fortran 轉(zhuǎn)換為C 進行了擴展。 擴展版本有151 個內(nèi)核。 本次研究選擇這個基準(zhǔn)測試是因為其中提供了一些相對簡單的循環(huán)的集合，這些循環(huán)可以在許多科學(xué)用途的C 語言代碼中找到，同時TSVC 中的循環(huán)能夠根據(jù)其測試的用途進行適當(dāng)分類，對結(jié)果分析有很大的幫助。 通過腳本將TSVC中每個函數(shù)提取出去生成獨立的C＋＋文件；TSVC的數(shù)據(jù)量對強化學(xué)習(xí)來說還是偏小了，本實驗將會在TSVC 原有的循環(huán)上進行增加，手段主要是改變變量名稱使得數(shù)據(jù)的總量倍增；同時，對不同分類的TSVC 函數(shù)進行隨機抽取用來構(gòu)建測試集。 構(gòu)建訓(xùn)練集和其中一個測試集如圖3 所示。

圖3 TSVC 訓(xùn)練集和測試集構(gòu)建流程Fig. 3 TSVC training and testing set construction process

為了測試本文構(gòu)建的強化學(xué)習(xí)循環(huán)展開化器能否很好地推廣到新代碼，并且測試其對于循環(huán)復(fù)雜程度的敏感程度，選擇使用Ameer 等學(xué)者［10］構(gòu)建的循環(huán)結(jié)構(gòu)簡單的數(shù)據(jù)集Tests、分離出的TSVC 內(nèi)核和PloyBench［11］三個測試集進行推廣性能的測試。這3 個測試集中，Tests 的循環(huán)結(jié)構(gòu)最簡單，其次是TSVC，PloyBench 的循環(huán)結(jié)構(gòu)最為復(fù)雜其主要由循環(huán)構(gòu)成，即由多個benchmark 組成。 本次研究選用其中的2mm、bicg、atax、gemm、gemver、choleshy 來進行測試。

3 實驗結(jié)果

強化學(xué)習(xí)循環(huán)展開器、LLVM 基準(zhǔn)模型和暴力搜索的性能對比如圖4 所示。 在Tests 基準(zhǔn)上本文提出的循環(huán)展開器比LLVM 基準(zhǔn)模型有1.19x 的性能提升；在TSVC 基準(zhǔn)上，有1.13x 的性能提升；在Polybench 基準(zhǔn)上，有1.09x 的性能提升。 可以看到，隨著循環(huán)復(fù)雜度的提升，循環(huán)展開器的性能也隨之下降，可見在面對復(fù)雜循環(huán)時本文提出的循環(huán)展開器還有進一步的優(yōu)化空間。

圖4 歸一化后強化學(xué)習(xí)循環(huán)展開器、LLVM 基準(zhǔn)模型和暴力搜索的性能比Fig. 4 The performance ratio of reinforcement learning loop unroller，LLVM benchmark model and brute force search after normalization

雖然本文提出的循環(huán)展開器性能上要弱于暴力搜索的結(jié)果，但是在運行速度上大幅領(lǐng)先暴力搜索，如圖5 所示：在Tests 基準(zhǔn)上，本文提出的循環(huán)展開器比暴力搜索快5.97x；在TSVC 基準(zhǔn)上，快6.23x 的性能提升；在Polybench 基準(zhǔn)上快5.61x 的性能提升。

圖5 歸一化后強化學(xué)習(xí)向量器和暴力搜索編譯所需時間比Fig. 5 Time ratio between reinforcement learning vector machineand brute force search compilation after normalization

4 結(jié)束語

針對編譯開發(fā)成本急需縮減的問題，本文提出了一個基于強化學(xué)習(xí)的循環(huán)展開器來對編譯開發(fā)中循環(huán)展開進行自動化。 通過代碼嵌入的方式來表征循環(huán)代碼，然后使用強化學(xué)習(xí)模型來對展開因子進行預(yù)測以達到輸入新代碼能快速自動地進行循環(huán)展開的目的。 實驗表明，本文提出的循環(huán)展開器相比于LLVM －O3 的基準(zhǔn)模型有更好的效果，同時有比暴力搜索更快的編譯速度。 但根據(jù)本文選取的3 個測試集的實驗結(jié)果，本文提出的向量化器在復(fù)雜場景中會效果下降，在未來的工作中會考慮更多復(fù)雜應(yīng)用場景，將謂詞加入到控制流的識別當(dāng)中、將多面體模型和強化學(xué)習(xí)進行結(jié)合以提升循環(huán)展開器的性能。