英偉達用AI設計GPU算術電路，面積比最先進EDA減少25%，速度更快、更有效率

大量的算術電路陣列為英偉達 GPU 提供了動力，以實現前所未有的 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）和電線構成。理想的電路應具有以下屬性：

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

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

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

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

研究者使用電路產生器將前綴圖轉換為一個帶有電線和邏輯閘的電路。接下來，這些產生的電路透過一個物理綜合工具來優化，該工具使用閘尺寸、複製和緩衝器插入等物理綜合優化。

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

研究者將算術電路設計視為一項強化學習（RL）任務，其中訓練一個智能體優化算術電路的面積和延遲屬性。對於前綴電路，他們設計了一個環境，​​其中RL智能體可以添加或刪除前綴圖中的節點，然後執行以下步驟：

1. ##前綴圖被規範化以始終保持正確的前綴和計算；
2. 從規範化的前綴圖中產生電路；
3. 使用物理綜合工具對電路進行物理綜合最佳化;
4. 測量電路的面積和延遲特性。

在下一個動圖中，RL智能體透過新增或刪除節點來一步一步地建立前綴圖。在每一步上，此智能體得到的獎勵是對應電路面積和延遲的改進。

原圖為可互動版本。

完全卷積Q學習智能體

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

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

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

Raptor進行分散式訓練

PrefixRL需要大量運算，在實體模擬中，每個GPU需要256個CPU，而訓練64b任務則需要超過32000個GPU小時。此次英偉達開發了一個內分散式強化學習平台Raptor，該平台充分利用了英偉達硬體優勢，從而可以進行這種工業級的強化學習（如下圖4）。

Raptor 能夠提高訓練模型的可擴展性和訓練速度，例如作業排程、自訂網路和 GPU 感知資料結構。在 PrefixRL 的上下文中，Raptor 使得跨 CPU、GPU 和 Spot 執行個體的混合分配成為可能。這個強化學習應用程式中的網路是多種多樣的，並且受益於以下幾點：

• #Raptor 在NCCL 之間切換以進行點對點傳輸，從而將模型參數直接從學習器GPU 傳送到推理GPU；
• Redis 用於非同步和較小的訊息，例如獎勵或統計資訊；
• 對於JIT 編譯的RPC，用於處理大容量和低延遲的請求，例如上傳經驗資料。

最後，Raptor 提供了GPU 感知資料結構，例如具有多執行緒服務的重播緩衝器，以接收來自多個worker的經驗，並行批次資料並將其預先載入到GPU 上。

下圖 4 顯示PrefixRL框架支援並發訓練和資料收集，並利用 NCCL 有效地向參與者（下圖中的actor）發送最新參數。

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

獎勵計算

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

RL 環境中的物理綜合最佳化可以產生各種解決方案來權衡面積和延遲。研究者使用與訓練特定智能體相同的權衡權重來驅動物理綜合工具。

在獎勵計算的循環中執行物理綜合優化具有以下優點：

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

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

研究者將獎勵計算與狀態更新解耦，因為智能體只需要當前的前綴圖狀態就可以採取行動，而不需要電路合成或先前的獎勵。由於Raptor，他們可以將冗長的獎勵運算轉移到CPU worker池中並行執行物理綜合，而actor智能體無需等待就能在環境中執行。

當CPU worker回傳獎勵時，轉換就可以嵌入重播緩衝器。綜合獎勵會被緩存，以避免再次遇到某個狀態時進行冗餘計算。

結果及展望

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

最好的 PrefixRL 加法器在相同延遲下實現的面積比 EDA 工具加法器低25%。這些在物理綜合最佳化後映射到Pareto最優加法器電路的前綴圖具有不規則的結構。

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

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

##據了解，這是第一個使用深度強化學習智能體來設計算術電路的方法。英偉達構思了一種藍圖：希望這種方法可以將 AI 應用於現實世界電路設計問題，建構動作空間、狀態表示、RL 智能體模型、針對多個競爭目標進行最佳化，以及克服緩慢的獎勵運算過程。

以上是英偉達用AI設計GPU算術電路，面積比最先進EDA減少25%，速度更快、更有效率的詳細內容。更多資訊請關注PHP中文網其他相關文章！