大量的算術電路陣列為英偉達 GPU 提供了動力，以實現(xiàn)前所未有的 AI、高性能計算和計算機圖形加速。因此，改進這些算術電路的設計對于提升 GPU 性能和效率而言至關重要。如果 AI 學習設計這些電路會怎么樣呢？在近期英偉達的論文《PrefixRL: Optimization of Parallel Prefix Circuits using Deep Reinforcement Learning》中，研究者證明了AI不僅可以從頭開始設計這些電路，而且AI設計的電路比最先進電子設計自動化（EDA）工具設計的電路更小、更快。

論文地址：https://arxiv.org/pdf/2205.07000.pdf

最新的英偉達Hopper GPU架構中擁有近13000個AI設計的電路實例。下圖 1 左 PrefixRL AI設計的64b加法器電路比圖1右最先進EDA工具設計的電路小25%。

?

電路設計概覽

計算機芯片中的算術電路是由邏輯門網絡（如NAND、NOR和XOR）和電線構成。理想的電路應具有以下屬性：

• 小：更小的面積，更多電路可以封裝在芯片上；
• 快：更低的延遲，提高芯片的性能；
• 更低功耗。

在英偉達的這項研究中，研究者關注電路面積和延遲。他們發(fā)現(xiàn)，功耗與感興趣電路的面積密切相關。電路面積和延遲往往是相互競爭的屬性，因此希望找到有效權衡這些屬性的設計的帕累托邊界。簡言之，研究者希望每次延遲時電路面積是最小的。

因此，在PrefixRL中，研究者專注于一類流行的算術電路——并行前綴電路。GPU中的各種重要電路如加速器、增量器和編碼器等都是前綴電路，它們可以在更高級別上被定為為前綴圖。

那么問題來了：AI智能體能設計出好的前綴圖嗎？所有前綴圖的狀態(tài)空間是很大的O(2^n^n)，無法使用蠻力方法進行探索。下圖2為具有4b電路實例的PrefixRL的一次迭代。

研究者使用電路生成器將前綴圖轉換為一個帶有電線和邏輯門的電路。接下來，這些生成的電路通過一個物理綜合工具來優(yōu)化，該工具使用門尺寸、復制和緩沖器插入等物理綜合優(yōu)化。

由于這些物理綜合優(yōu)化，最終的電路屬性（延遲、面積和功率）不會直接從原始前綴圖屬性（如電平和節(jié)點數）轉換而來。這就是為什么AI智能體學習設計前綴圖但又要對從前綴圖中生成的最終電路的屬性進行優(yōu)化。

研究者將算術電路設計視為一項強化學習（RL）任務，其中訓練一個智能體優(yōu)化算術電路的面積和延遲屬性。對于前綴電路，他們設計了一個環(huán)境，其中RL智能體可以添加或刪除前綴圖中的節(jié)點，然后執(zhí)行如下步驟：

1. 前綴圖被規(guī)范化以始終保持正確的前綴和計算；
2. 從規(guī)范化的前綴圖中生成電路；
3. 使用物理綜合工具對電路進行物理綜合優(yōu)化；
4. 測量電路的面積和延遲特性。

在如下動圖中，RL智能體通過添加或刪除節(jié)點來一步步地構建前綴圖。在每一步上，該智能體得到的獎勵是對應電路面積和延遲的改進。

原圖為可交互版本。

完全卷積Q學習智能體

研究者采用Q學習（Q-learning）算法來訓練智能體電路設計。如下圖3所示，他們將前綴圖分解成網格表示，其中網格中的每個元素唯一地映射到前綴節(jié)點。這種網格表示用于 Q 網絡的輸入和輸出。輸入網格中的每個元素表示節(jié)點是否存在。輸出網格中的每個元素代表添加或刪除節(jié)點的 Q 值。

研究者采用完全卷積神經網絡架構，因為Q學習智能體的輸入和輸出都是網格表示。智能體分別預測面積和延遲屬性的 Q 值，因為面積和延遲的獎勵在訓練期間是單獨可觀察的。

圖3：4b前綴圖表示（左）以及完全卷積Q學習智能體架構（右）。

Raptor進行分布式訓練

PrefixRL需要大量計算，在物理模擬中，每個GPU需要256個CPU，而訓練64b任務則需要超過32000個GPU小時。此次英偉達開發(fā)了一個內分布式強化學習平臺Raptor，該平臺充分利用了英偉達硬件優(yōu)勢，從而可以進行這種工業(yè)級別的強化學習（如下圖4）。

Raptor 能夠提高訓練模型的可擴展性和訓練速度，例如作業(yè)調度、自定義網絡和 GPU 感知數據結構。在 PrefixRL 的上下文中，Raptor 使得跨 CPU、GPU 和 Spot 實例的混合分配成為可能。這個強化學習應用程序中的網絡是多種多樣的，并且受益于以下幾點：

• Raptor 在 NCCL 之間切換以進行點對點傳輸，從而將模型參數直接從學習器 GPU 傳輸到推理 GPU；
• Redis 用于異步和較小的消息，例如獎勵或統(tǒng)計信息；
• 對于 JIT 編譯的 RPC，用于處理大容量和低延遲的請求，例如上傳經驗數據。

最后，Raptor 提供了 GPU 感知數據結構，例如具有多線程服務的重放緩沖器，以接收來自多個worker的經驗，并行批處理數據并將其預先載入到 GPU 上。

下圖 4 顯示PrefixRL框架支持并發(fā)訓練和數據收集，并利用 NCCL 有效地向參與者（下圖中的actor）發(fā)送最新參數。

圖4：研究者采用Raptor進行解耦并行訓練和獎勵計算，以克服電路合成延遲。

獎勵計算

研究者采用權衡權重 w （范圍為[0,1]）來組合區(qū)域和延遲目標。他們訓練具有不同權重的各種智能體以獲得帕累托邊界，從而平衡面積、延遲之間的權衡。

RL 環(huán)境中的物理綜合優(yōu)化可以生成各種解決方案來權衡面積和延遲。研究者使用與訓練特定智能體相同的權衡權重來驅動物理綜合工具。

在獎勵計算的循環(huán)中執(zhí)行物理綜合優(yōu)化具有以下優(yōu)點：

• RL智能體學習直接優(yōu)化目標技術節(jié)點和庫的最終電路屬性；
• RL智能體在物理綜合過程中包含目標算法電路的周邊邏輯，從而共同優(yōu)化目標算法電路及其周邊邏輯的性能。

然而，進行物理綜合是一個緩慢的過程(64b加法器~35秒)，這可能大大減慢RL的訓練和探索。

研究者將獎勵計算與狀態(tài)更新解耦，因為智能體只需要當前的前綴圖狀態(tài)就可以采取行動，而不需要電路合成或之前的獎勵。得益于Raptor，他們可以將冗長的獎勵計算轉移到CPU worker池中并行執(zhí)行物理綜合，而actor智能體無需等待就能在環(huán)境中執(zhí)行。

當CPU worker返回獎勵時，轉換就可以嵌入重放緩沖器。綜合獎勵會被緩存，以避免再次遇到某個狀態(tài)時進行冗余計算。

結果及展望

下圖5展示了使用PrefixRL設計的64b加法器電路、以及來自最先進EDA工具的帕累托支配加法器電路的面積和延遲情況。

最好的 PrefixRL 加法器在相同延遲下實現(xiàn)的面積比 EDA 工具加法器低25%。這些在物理綜合優(yōu)化后映射到Pareto最優(yōu)加法器電路的前綴圖具有不規(guī)則的結構。

?圖5：PrefixRL設計的算術電路比最先進的EDA工具設計的電路更小和更快。

(左)電路架構；(右)對應的64b加法器電路特性圖

據了解，這是首個使用深度強化學習智能體來設計算術電路的方法。英偉達構想了一種藍圖：希望這種方法可以將 AI 應用于現(xiàn)實世界電路設計問題，構建動作空間、狀態(tài)表示、RL 智能體模型、針對多個競爭目標進行優(yōu)化，以及克服緩慢的獎勵計算過程。