6小時，Jeff Dean帶AI設計晶片方案登Nature，谷歌第四代TPU已用

機器之心報道

編輯：杜偉、陳萍

將晶片的佈局規劃看作一個深度強化學習問題，谷歌大腦團隊希望用 AI 來提升晶片設計效率。基於 AI 的最新設計方案可以在數小時內完成人類設計師耗費數月才能完成的晶片佈局，這將有可能引領一場新的晶片效率革命。

2020 年 4 月，包括 Google AI 負責人 Jeff Dean 在內的谷歌大腦研究者描述了一種基於 AI 的晶片設計方法，該方法可以從過往經驗中學習並隨時間推移不斷改進，從而能夠更好地生成不可見（unseen）元件的架構。據他們表示，這種基於 AI 的方法平均可以在 6 小時內完成設計，這要比人類專家所需要的數週時間快得多。

左為人類設計的微晶片佈局，右為機器學習系統設計的微晶片佈局。圖源：Nature

近日，谷歌大腦團隊聯合斯坦福大學的研究者對這一基於 AI 的晶片設計方法進行了改進，並將其應用於不久前 Google I/O 2021 大會上正式釋出的、下一代張量處理單元（TPU v4）加速器的產品中。谷歌此前表示，TPUv4 可以在目標檢測、影象分類、自然語言處理、機器翻譯和推薦基準等工作負載上優於上一代 TPU 產品。

相關論文研究已經在 Nature 上發表，Jeff Dean 為核心作者之一。據介紹，在不到六小時的時間內，谷歌 AI 晶片設計方法自動生成的晶片佈局在功耗、效能和芯片面積等所有關鍵指標上都優於或媲美人類，而工程師需要耗費數月的艱苦努力才能達到類似效果。

論文地址：https://www.nature.com/articles/s41586-021-03544-w

這項基於強化學習的快速晶片設計方法對於資金緊張的初創企業大有裨益，如果谷歌公開相關技術的話，這些初創企業可以開發自己的 AI 和其他專用晶片。並且，這種方法有助於縮短晶片設計週期，從而使得硬體可以更好地適應快速發展的技術研究。

技術詳解

晶片佈局是設計計算機晶片物理佈局的一項重要工程任務。在電子設計自動化（EDA）出現之前，設計人員必須手工完成積體電路的設計、佈線等工作，到了 1970 年代中期，開發人員嘗試將整個設計過程自動化。此後，第一個電路佈局佈線工具研發成功，設計自動化研討會（Design Automation Conference）在這一時期被創立。電子設計自動化發展的下一個重要階段以卡弗爾 · 米德（Carver Mead）和琳 · 康維於 1980 年發表的論文《超大規模積體電路系統導論》，提出了通過程式語言來進行晶片設計的新思想。從 1981 年開始，電子設計自動化逐漸開始商業化。

儘管歷經了 50 年的相關研究，晶片佈局仍與自動化背道而馳，需要物理設計工程師數月的艱苦努力才能生產出可製造的佈局。基於此，谷歌研究者提出了一種用於晶片佈局設計的深度強化學習（RL）方法。就其效果而言，在不到六小時的時間內，谷歌設計方法自動生成的晶片佈局在功耗、效能和芯片面積等所有關鍵指標上都優於或媲美人類。

具體而言，為了實現這一目標，研究者將晶片佈局作為一個強化學習問題，開發了一種基於邊緣、能夠學習晶片豐富且可遷移表示的圖卷積神經網路架構。這種方法能夠更好地利用過往的經驗，從而更好更快地解決問題的新例項，使得晶片設計由比任何人類設計師具備更多經驗的人工智慧體執行。

此外，這種方法可被用於設計谷歌下一代人工智慧加速器，並且有可能為它們節省數千小時的人力。研究者相信，更強大的 AI 設計的硬體將推動人工智慧領域的進步，並在這兩個領域之間建立一種共生關係。

設計域 - 自適應策略

為晶片佈局規劃開發域 - 自適應策略是非常具有挑戰性的，因為該問題類似於具有不同棋子、棋盤和獲勝條件的遊戲，元件是「棋子」（例如，網表拓撲、宏計數、宏大小和縱橫比）、放置元件的畫布是「棋盤」（畫布大小和長寬比）、贏的條件（不同的評估指標或不同的密度和路由擁塞約束的相對重要性）。即使是遊戲的一個例項（將一個特定的網表放到一個特定的畫布上）也有一個巨大的狀態 - 動作空間，對全局形成影響。為了應對這個挑戰，研究者首先集中學習狀態空間的豐富表示。

研究者訓練了一個神經網路架構，能夠預測新網表放置的獎勵，最終目標是將此架構用作整個策略的編碼層。

為了訓練這種監督模型，需要大量的晶片放置資料集及其相應的獎勵標籤。因此，研究者創建了一個包含 10000 個晶片放置的資料集，其中輸入是與給定放置相關的狀態，標籤是該放置的獎勵。為了準確預測獎勵標籤，將其泛化到未看到的資料中，研究者提出了一種基於邊緣的圖神經網路架構，即 Edge-GNN（基於邊緣圖神經網路）。該網路的作用是將網表嵌入，將節點類型和連通性的資訊提取到低維向量表示中，以用於下游任務。

監督模型通過迴歸進行訓練，以最小化均方損失的加權和。監督任務使研究者能夠找到在網表中推廣獎勵預測所需的特徵和架構。為了將 Edge-GNN 合併到 RL 策略網路中，該研究移除了預測層，然後將其用作策略網路的編碼器，如圖所示。

策略網路和價值網路體系架構。

訓練方法和訓練方案。

谷歌團隊的系統從一個空晶片開始，按順序放置元件，直到完成網表。為了指導系統選擇先放置哪些元件，元件按大小遞減進行排序，研究人員先放置大的元件，這樣可以減少之後無法放置的可能性。

開源 RISC-V 處理器 Ariane 的宏佈局隨著訓練程序的變化情況。左邊：從頭開始訓練；右邊：正在為這個晶片調整預訓練策略。每個矩形代表一個單獨的宏位置。

自適應結果

下圖 3 中，研究人員比較了使用預訓練策略生成的放置質量，以及通過從頭開始訓練策略生成的放置質量，訓練資料集由 TPU 塊和開源 Ariane RISC-V CPU 塊組成，在每一個實驗中，都對除目標塊外的所有塊預訓練策略。

研究人員展示了零樣本模式的結果，以及在特定設計上對預訓練策略進行 2 小時和 12 小時微調後的結果。

從頭開始訓練的策略收斂時間要長得多，即使在 24 小時後（根據獎勵函數評估），結果也比微調策略在 12h 內達到的效果還要差。這表明，在預訓練期間接觸許多不同的設計可以更快地為新的看不見的塊生成更高質量的放置方案。

圖 4 顯示了從頭開始訓練與來自 Ariane RISC-V CPU30 的預訓練策略網路的訓練的收斂圖。結果顯示預訓練策略不僅具有較低的放置成本，而且能比從頭開始訓練的策略收斂速度快 30 小時以上。

下圖 5 展示了規模更大的訓練集對效能的影響。研究者依次將訓練集從 2 個 TPU 塊增至 5 個，最後增至 20 個，策略網路在零樣本和經過相同小時數的微調後均能生成更好的晶片佈局。這表明，在將策略網路應用於更多不同類型的晶片設計過程中，它不易於出現過擬合，並且能夠更好地泛化至新的未知晶片設計。

與基線方法對比

研究者將該方法與當前 SOTA 方法以及人類設計師團隊完成的上一代 TPU 的產品設計進行了對比。為了公平起見，研究者確保所有方法使用相同的實驗設定，包括相同的輸入和 EDA 工具設定，並使用在最大資料集（20 個 TPU 塊）上預先訓練的 AI 策略，接著在 5 個目標不可見塊上微調（時間少於 6 小時）。

結果如下表 1 所示，表明基於 AI 的方法在生成滿足設計要求的高質量晶片佈局方面是有效的，在面積、功率和線長方面均優於或媲美人類專家手動設計的效果。

為晶片設計過程的完全自動化奠定基礎

谷歌稱其系統泛化和生成高質量解決方案的能力具有重大影響，為晶片設計過程的早期優化提供了機會。以前，大規模的架構探索是不可能的，因為評估一個給定的架構候選需要花費數月的努力。然而，谷歌團隊指出，修改晶片設計可能會對效能產生巨大影響，並可能為晶片設計過程的完全自動化奠定基礎。

此外，由於谷歌團隊的系統只是學習將一個圖的節點對映到一組資源上，因此它可能適用於包括城市規劃、疫苗測試和分發以及大腦皮層對映在內的一系列應用。

研究人員在論文中寫道：「（雖然）我們的方法已經在生產中被用於設計下一代谷歌 TPU…… 我們相信，（它）可以應用於晶片設計以外的有影響力的佈局問題。」