Jeff Dean領銜AI助力芯片設計效率革命！6小時內完成布局設計，新一代TPU已用上

大數據文摘授權轉載自學術頭條

撰稿：庫珀 

設計一款高端芯片有多難？想象一下要將上百億個晶體管集成到指甲蓋大小的區域上，還要進行各類功能區域規劃，這是一項多么微觀且浩瀚的工程。

如今，最基本的芯片晶體管部件都是微米、納米級尺度，如何鋪設和利用這幾百億個晶體管，如何通過模塊設計組合迸發出一款芯片的極致性能，如何滿足日益增長的終端需求，這是如今芯片設計工程師們面臨的嚴苛挑戰，而且如此復雜的設計，不能有紕漏缺陷，一旦投入量產無法補救，必須復盤再來，而如果一款芯片設計方案足夠好，甚至能彌補一定程度的制程差距。

現在，隨著 AI 技術深入到芯片設計領域，一場效率革命可能要再次席卷而來。

來自由 Jeff Dean 領銜的谷歌大腦團隊以及斯坦福大學計算機科學系的科學家們，在一項聯合研究中證明，一種基于深度強化學習的芯片布局規劃方法，能夠生成可行的芯片設計方案。

在不到 6 小時的時間內，由該方法自動生成的芯片平面圖在所有關鍵指標上（包括功耗、性能和芯片面積等參數）都優于或與人類專家生成的設計圖效果相當，而人類工程師往往需要數月的緊張努力才能達到如此效果。

圖｜人類設計的微芯片平面圖（a）與機器學習系統設計（b）的不同（來源：Nature） 

據論文描述，如今這種方法已經被用于下一代 Google 張量處理單元（TPU）加速器的產品中，有望為今后的每一代計算機芯片迭代節省數千小時的人力，相關論文于 6 月 10 日發表在頂級科學期刊 Nature 上，谷歌 AI 掌門人 Jeff Dean 是核心作者之一。

（來源：Nature）

芯片設計工程師的AI助理

事實上，科學家們一直在探索讓芯片設計過程更高效的路徑。

重大的進步出現在 1980 年代，電子設計自動化（EDA）技術誕生，開發人員能將整個設計過程自動化，利用計算機輔助設計（CAD）軟件，完成超大規模集成電路（VLSI）的功能設計、綜合、驗證、物理設計等復雜流程，EDA 的不斷完善架起了芯片設計與制造之間的橋梁，在半導體產業鏈中有著舉足輕重的作用。

盡管已經經過了 50 年的相關研究，但芯片布局規劃仍無法實現自動化，依然需要物理設計工程師數月精心布置才能完成可制造的芯片布局。

具體而言，一個芯片往往會被分成幾十個模塊，每個模塊都有單獨的功能屬性，例如內存子系統、計算單元或控制邏輯系統，這些模塊可以用網表（也稱連線表）、電路組件的超圖和標準單元（邏輯門，例如 NAND、NOR 和 XOR）來描述，所有這些模塊再通過導線連接起來。

芯片布局規劃則涉及將網表放置在芯片畫布（二維網格）上，以便優化性能指標（例如功耗、定時、面積和線長等），同時遵守密度和路由擁塞的硬性約束條件。

其實，自 20 世紀 60 年代以來已有許多芯片布局規劃方法被提出，大概分為三類：

基于分區的方法（partitioning-based methods）；

隨機/爬山方法（stochastic/hill-climbing approaches）；

解析解算器（analytic solvers）。

但是，這些方法目前都無法達到人類專家級水平的性能，因為各有各的局限。

基于分區的方法為了擴展到更大的網表而犧牲了全局解決方案的質量，而一個糟糕的早期分區可能會導致一個無法解決的最終結果；

隨機/爬山方法具有較低的設計收斂率，不能擴展到具有數百萬乃至數十億節點的現代芯片網表；

解析解算器是一種主流方法，但它只能優化可微損失函數，這意味著它們不能有效地優化關鍵指標，如路由擁塞或時序沖突等。

由于這些先前方法的局限性，人類物理設計師必須使用商業 EDA 工具迭代數月，將芯片網表的寄存器傳輸級（RTL）描述作為輸入，在芯片畫布上進行手動放置，并等待長達 72 小時的時間，讓 EDA 工具來評估該位置是否最佳。

在這個反饋的基礎上，人類設計師要么得出設計標準達到目標的結論，生成一個更新的層級平面圖進行評估，要么得向上游 RTL 設計師提供反饋，然后上游 RTL 設計師修改低級代碼以使放置任務更容易。

為了解決這類復雜的芯片布局規劃問題，研究人員開發了一種強化學習方法，該方法能夠跨芯片進行推廣，這意味著它可以從經驗中不斷學習，在放置新芯片模塊時變得更好更快，從而使芯片設計師能夠得到 AI 的協助，且這個 AI 助理的“經驗”比任何人都豐富。

但培訓跨芯片的 AI 布局策略非常具有挑戰性，因為它需要學習優化所有可能的芯片網表到所有可能的芯片畫布上的布局。

這類似于一個棋盤游戲，元件是“棋子”（例如網表拓撲、宏計數、宏大小和縱橫比）、放置元件的畫布是“棋盤”（不同的芯片畫布大小和縱橫比）和獲勝條件（不同的評估指標或不同的密度和路由擁塞約束的相對重要性）。即使是這個游戲的一個實例（將一個特定的網表放到一個特定的畫布上）也有一個巨大的狀態-動作空間，對全局形成影響。

為了實現泛化，研究人員重點研究了芯片的可轉移表征學習，并將表征學習應用于預測芯片放置質量的有監督任務中。通過設計一種能夠準確預測各種網表及其位置的神經結構，研究人員能夠生成輸入網表的豐富特征嵌入，然后，使用這個結構作為編碼器，以實現轉移學習。

除了對芯片布局規劃的直接影響外，基于 AI 的芯片布局設計方案概括和快速生成高質量解決方案的能力也具有重大意義，為與芯片設計過程早期階段的協同優化創造了機會。

在更抽象的層次上，研究人員表示，這種形式的位置優化策略還可出現在廣泛的科學和工程應用中，例如硬件設計、城市規劃、疫苗測試和分發以及大腦皮層布局研究等。

因此，研究人員相信，這種 AI 布局優化方法策略可以應用于芯片設計以外更廣泛的有效布局問題。

設計域的自適應策略

為芯片布局規劃開發領域自適應策略非常具有挑戰性，因為每個芯片可以包含數千萬個邏輯門（稱為標準單元）以及數千個內存塊（稱為宏塊或宏），單元和宏塊通過數十公里的布線相互連接，以實現設計的功能，其中蘊藏著驚人的復雜性，工程師們在專業軟件工具的幫助下，才能控制一部分復雜性。

其中，標準單元和宏塊在芯片中的位置對設計結果至關重要，它們的位置決定了導線必須跨越的距離，從而影響元件之間的布線是否成功，以及邏輯門之間的信號傳輸速度。因此，隨著邏輯電路設計的發展，人類工程師必須反復調整宏塊的位置，每次迭代都需要手動執行，需要幾天或幾周的時間。

眾所周知，計算機行業是由摩爾定律驅動的，每一塊芯片的元件數量大約每兩年翻一番，相當于一個芯片上的元件數量每周增加約 1%。因此，無法自動完成布局規劃是一個迫切的問題，這不僅是因為相關的時間成本會增加，也限制了芯片開發計劃中可以探索的解決方案的數量。

如何讓 AI 代理成功地將宏塊一個接一個合理地放入芯片布局中？

具體而言，研究人員首先集中學習狀態空間的豐富表征，他們的直覺是，一個能夠完成芯片放置一般任務的 AI 策略，也應該能夠在推理時將與一個新的不可見芯片相關的狀態編碼成一個有意義的信號。

為此，研究人員首先訓練了一個神經網絡架構，能夠預測新網表放置時的回報，最終目標是將此架構用作整個策略的編碼層，訓練這個架構需要一個大的芯片位置數據集和相應的獎勵標簽，研究人員創建了一個包含 10000 個芯片放置的數據集。

另外，為了準確預測獎賞標簽并推廣到未知數據，研究人員又開發了一種神經網絡結構：Edge-GNN (基于邊緣的圖形神經網絡)，該網絡的作用是嵌入網表，將節點類型和連通性的信息提取到低維向量表示中，以用于下游任務。

圖｜策略和價值網絡架構（來源：該論文）

嵌入層會對有關網表鄰接、節點特征和要放置的當前宏的信息進行編碼，然后，策略和值網絡分別輸出可用網格單元上的概率分布和當前布局的預期回報估計。

圖｜整體方法和訓練方案概述（來源：該論文）

在每個訓練迭代中，強化學習代理一次放置一個宏（一種批量處理的稱謂），放置完所有宏后，將使用強制方法放置標準單元，在每次迭代結束時，獎勵被計算為近似線長、擁塞和密度的線性組合，并作為反饋提供給強化學習代理，以便為下一次迭代優化其參數。

域自適應結果

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

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

圖｜從無到有的訓練與不同時間量的微調，從表中可以看出，預先訓練的策略網絡始終優于從無到有訓練的策略網絡，證明了從離線訓練數據中學習的有效性（來源：該論文）

圖｜Ariane RISC-V CPU30 從零開始訓練與從預先訓練的策略網絡開始訓練的收斂圖（來源：該論文）

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

預先訓練的策略不僅具有較低的放置成本，而且能比從頭開始訓練的策略收斂速度快 30 小時以上。

接下來，研究人員探討了培訓數據對策略學習能力的影響，例如從更大的數據集中學習。TPU 芯片塊非常多樣化，他們仔細選擇了一系列具有代表性的功能塊（例如片上和片間網絡塊、計算核心、內存控制器、數據傳輸緩沖區和邏輯以及各種接口控制器）、飽和（宏的總面積與畫布的總面積之比，<30%，30%-60% 和 >60%）和宏計數（最多幾百個），小訓練集包含 2 個塊，中訓練集包含 5 個塊，大訓練集包含 20 個塊。

圖｜訓練前數據集大小的影響（來源：該論文）

當研究人員對更多的芯片塊進行預訓練時，能夠更快地為新的看不見的芯片塊生成更高質量的放置。當他們將訓練集從 2 個塊增加到 5 個塊，最后增加到 20 個塊時，策略網絡在微調后都會生成更好的放置位置。

這表明，將策略暴露于更多不同的芯片設計時，它變得不太容易過度擬合，并且更善于推廣到新的未知設計中。

研究人員也探討了 AI 芯片布局設計與基線方法的數據比較。在評估芯片布局圖的質量時，有幾個重要指標，它們相互權衡。沒有一個度量可以用來捕獲放置的總體質量，因此報告中的關鍵度量，例如包括總的線長、計時、路由擁塞（水平和垂直）、面積和功率等。

為了進行公平比較，研究人員確保所有方法都有相同的實驗設置，包括相同的輸入和相同的EDA工具設置，使用一個在最大的數據集（20 個 TPU 塊）上預先訓練的 AI 策略，然后在 5 個目標不可見塊上微調它，時間不超過 6 小時。

圖｜與基線的比較，AI 布局芯片的方法與最先進的方法，以及使用行業標準 EDA 工具手動放置進行比較。對于本表中的所有指標，越低越好（來源：該論文）

結果證明，利用 AI 進行芯片設計布局是有效的，不僅能生成高質量的模塊放置方案，也能滿足實際設計標準要求，在面積、功率和線長方面優于或不亞于人類專家手動放置的效果。

加利福尼亞大學圣迭戈分校電氣與計算機工程系專家 Andrew B. Kahng 對這項成果評論稱，對于芯片設計和設計自動化領域的長期實踐者來說，這項研究成果確實很神奇。

最重要的啟示可能是，該技術方案已經被納入谷歌下一代 AI 處理器的芯片設計中，這意味著這種解決方案足夠好、足夠成熟，可以在昂貴、尖端的硅片上批量復制，因此，未來可以期望半導體行業從業者加倍關注此類相關工作，并在整個芯片設計過程中尋求大量類似的應用落地。

參考鏈接：

https://www.nature.com/articles/s41586-021-03544-w

https://www.nature.com/articles/d41586-021-01515-9