英偉達用AI設(shè)計GPU算術(shù)電路，面積比最先進EDA減少25%，速度更快、更加高效
大量的算術(shù)電路陣列為英偉達 GPU 提供了動力，以實現(xiàn)前所未有的 AI、高性能計算和計算機圖形加速。因此，改進這些算術(shù)電路的設(shè)計對于提升 GPU 性能和效率而言至關(guān)重要。如果 AI 學(xué)習(xí)設(shè)計這些電路會怎么樣呢？在近期英偉達的論文《PrefixRL: Optimization of Parallel Prefix Circuits using Deep Reinforcement Learning》中，研究者證明了AI不僅可以從頭開始設(shè)計這些電路，而且AI設(shè)計的電路比最先進電子設(shè)計自動化（EDA）工具設(shè)計的電路更小、更快。
論文地址：https://arxiv.org/pdf/2205.07000.pdf
最新的英偉達Hopper GPU架構(gòu)中擁有近13000個AI設(shè)計的電路實例。下圖 1 左 PrefixRL AI設(shè)計的64b加法器電路比圖1右最先進EDA工具設(shè)計的電路小25%。
?
電路設(shè)計概覽
計算機芯片中的算術(shù)電路是由邏輯門網(wǎng)絡(luò)（如NAND、NOR和XOR）和電線構(gòu)成。理想的電路應(yīng)具有以下屬性：
- ?。焊〉拿娣e，更多電路可以封裝在芯片上；
- 快：更低的延遲，提高芯片的性能；
- 更低功耗。
在英偉達的這項研究中，研究者關(guān)注電路面積和延遲。他們發(fā)現(xiàn)，功耗與感興趣電路的面積密切相關(guān)。電路面積和延遲往往是相互競爭的屬性，因此希望找到有效權(quán)衡這些屬性的設(shè)計的帕累托邊界。簡言之，研究者希望每次延遲時電路面積是最小的。
因此，在PrefixRL中，研究者專注于一類流行的算術(shù)電路——并行前綴電路。GPU中的各種重要電路如加速器、增量器和編碼器等都是前綴電路，它們可以在更高級別上被定為為前綴圖。
那么問題來了：AI智能體能設(shè)計出好的前綴圖嗎？所有前綴圖的狀態(tài)空間是很大的O(2^n^n)，無法使用蠻力方法進行探索。下圖2為具有4b電路實例的PrefixRL的一次迭代。
研究者使用電路生成器將前綴圖轉(zhuǎn)換為一個帶有電線和邏輯門的電路。接下來，這些生成的電路通過一個物理綜合工具來優(yōu)化，該工具使用門尺寸、復(fù)制和緩沖器插入等物理綜合優(yōu)化。
由于這些物理綜合優(yōu)化，最終的電路屬性（延遲、面積和功率）不會直接從原始前綴圖屬性（如電平和節(jié)點數(shù)）轉(zhuǎn)換而來。這就是為什么AI智能體學(xué)習(xí)設(shè)計前綴圖但又要對從前綴圖中生成的最終電路的屬性進行優(yōu)化。
研究者將算術(shù)電路設(shè)計視為一項強化學(xué)習(xí)（RL）任務(wù)，其中訓(xùn)練一個智能體優(yōu)化算術(shù)電路的面積和延遲屬性。對于前綴電路，他們設(shè)計了一個環(huán)境，其中RL智能體可以添加或刪除前綴圖中的節(jié)點，然后執(zhí)行如下步驟：
- 前綴圖被規(guī)范化以始終保持正確的前綴和計算；
- 從規(guī)范化的前綴圖中生成電路；
- 使用物理綜合工具對電路進行物理綜合優(yōu)化；
- 測量電路的面積和延遲特性。
在如下動圖中，RL智能體通過添加或刪除節(jié)點來一步步地構(gòu)建前綴圖。在每一步上，該智能體得到的獎勵是對應(yīng)電路面積和延遲的改進。
原圖為可交互版本。
完全卷積Q學(xué)習(xí)智能體
研究者采用Q學(xué)習(xí)（Q-learning）算法來訓(xùn)練智能體電路設(shè)計。如下圖3所示，他們將前綴圖分解成網(wǎng)格表示，其中網(wǎng)格中的每個元素唯一地映射到前綴節(jié)點。這種網(wǎng)格表示用于 Q 網(wǎng)絡(luò)的輸入和輸出。輸入網(wǎng)格中的每個元素表示節(jié)點是否存在。輸出網(wǎng)格中的每個元素代表添加或刪除節(jié)點的 Q 值。
研究者采用完全卷積神經(jīng)網(wǎng)絡(luò)架構(gòu)，因為Q學(xué)習(xí)智能體的輸入和輸出都是網(wǎng)格表示。智能體分別預(yù)測面積和延遲屬性的 Q 值，因為面積和延遲的獎勵在訓(xùn)練期間是單獨可觀察的。
圖3：4b前綴圖表示（左）以及完全卷積Q學(xué)習(xí)智能體架構(gòu)（右）。
Raptor進行分布式訓(xùn)練
PrefixRL需要大量計算，在物理模擬中，每個GPU需要256個CPU，而訓(xùn)練64b任務(wù)則需要超過32000個GPU小時。此次英偉達開發(fā)了一個內(nèi)分布式強化學(xué)習(xí)平臺Raptor，該平臺充分利用了英偉達硬件優(yōu)勢，從而可以進行這種工業(yè)級別的強化學(xué)習(xí)（如下圖4）。
Raptor 能夠提高訓(xùn)練模型的可擴展性和訓(xùn)練速度，例如作業(yè)調(diào)度、自定義網(wǎng)絡(luò)和 GPU 感知數(shù)據(jù)結(jié)構(gòu)。在 PrefixRL 的上下文中，Raptor 使得跨 CPU、GPU 和 Spot 實例的混合分配成為可能。這個強化學(xué)習(xí)應(yīng)用程序中的網(wǎng)絡(luò)是多種多樣的，并且受益于以下幾點：
- Raptor 在 NCCL 之間切換以進行點對點傳輸，從而將模型參數(shù)直接從學(xué)習(xí)器 GPU 傳輸?shù)酵评?GPU；
- Redis 用于異步和較小的消息，例如獎勵或統(tǒng)計信息；
- 對于 JIT 編譯的 RPC，用于處理大容量和低延遲的請求，例如上傳經(jīng)驗數(shù)據(jù)。
最后，Raptor 提供了 GPU 感知數(shù)據(jù)結(jié)構(gòu)，例如具有多線程服務(wù)的重放緩沖器，以接收來自多個worker的經(jīng)驗，并行批處理數(shù)據(jù)并將其預(yù)先載入到 GPU 上。
下圖 4 顯示PrefixRL框架支持并發(fā)訓(xùn)練和數(shù)據(jù)收集，并利用 NCCL 有效地向參與者（下圖中的actor）發(fā)送最新參數(shù)。
圖4：研究者采用Raptor進行解耦并行訓(xùn)練和獎勵計算，以克服電路合成延遲。
獎勵計算
研究者采用權(quán)衡權(quán)重 w （范圍為[0,1]）來組合區(qū)域和延遲目標。他們訓(xùn)練具有不同權(quán)重的各種智能體以獲得帕累托邊界，從而平衡面積、延遲之間的權(quán)衡。
RL 環(huán)境中的物理綜合優(yōu)化可以生成各種解決方案來權(quán)衡面積和延遲。研究者使用與訓(xùn)練特定智能體相同的權(quán)衡權(quán)重來驅(qū)動物理綜合工具。
在獎勵計算的循環(huán)中執(zhí)行物理綜合優(yōu)化具有以下優(yōu)點：
- RL智能體學(xué)習(xí)直接優(yōu)化目標技術(shù)節(jié)點和庫的最終電路屬性；
- RL智能體在物理綜合過程中包含目標算法電路的周邊邏輯，從而共同優(yōu)化目標算法電路及其周邊邏輯的性能。
然而，進行物理綜合是一個緩慢的過程(64b加法器~35秒)，這可能大大減慢RL的訓(xùn)練和探索。
研究者將獎勵計算與狀態(tài)更新解耦，因為智能體只需要當前的前綴圖狀態(tài)就可以采取行動，而不需要電路合成或之前的獎勵。得益于Raptor，他們可以將冗長的獎勵計算轉(zhuǎn)移到CPU worker池中并行執(zhí)行物理綜合，而actor智能體無需等待就能在環(huán)境中執(zhí)行。
當CPU worker返回獎勵時，轉(zhuǎn)換就可以嵌入重放緩沖器。綜合獎勵會被緩存，以避免再次遇到某個狀態(tài)時進行冗余計算。
結(jié)果及展望
下圖5展示了使用PrefixRL設(shè)計的64b加法器電路、以及來自最先進EDA工具的帕累托支配加法器電路的面積和延遲情況。
最好的 PrefixRL 加法器在相同延遲下實現(xiàn)的面積比 EDA 工具加法器低25%。這些在物理綜合優(yōu)化后映射到Pareto最優(yōu)加法器電路的前綴圖具有不規(guī)則的結(jié)構(gòu)。
?圖5：PrefixRL設(shè)計的算術(shù)電路比最先進的EDA工具設(shè)計的電路更小和更快。
(左)電路架構(gòu)；(右)對應(yīng)的64b加法器電路特性圖
據(jù)了解，這是首個使用深度強化學(xué)習(xí)智能體來設(shè)計算術(shù)電路的方法。英偉達構(gòu)想了一種藍圖：希望這種方法可以將 AI 應(yīng)用于現(xiàn)實世界電路設(shè)計問題，構(gòu)建動作空間、狀態(tài)表示、RL 智能體模型、針對多個競爭目標進行優(yōu)化，以及克服緩慢的獎勵計算過程。
