Jeff Dean領(lǐng)銜AI助力芯片設(shè)計效率革命！6小時內(nèi)完成布局設(shè)計，新一代TPU已用上

大數(shù)據(jù)文摘授權(quán)轉(zhuǎn)載自學(xué)術(shù)頭條

撰稿：庫珀 

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

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

現(xiàn)在，隨著 AI 技術(shù)深入到芯片設(shè)計領(lǐng)域，一場效率革命可能要再次席卷而來。

來自由 Jeff Dean 領(lǐng)銜的谷歌大腦團隊以及斯坦福大學(xué)計算機科學(xué)系的科學(xué)家們，在一項聯(lián)合研究中證明，一種基于深度強化學(xué)習(xí)的芯片布局規(guī)劃方法，能夠生成可行的芯片設(shè)計方案。

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

圖｜人類設(shè)計的微芯片平面圖（a）與機器學(xué)習(xí)系統(tǒng)設(shè)計（b）的不同（來源：Nature） 

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

（來源：Nature）

芯片設(shè)計工程師的AI助理

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

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

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

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

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

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

基于分區(qū)的方法（partitioning-based methods）；

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

解析解算器（analytic solvers）。

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

基于分區(qū)的方法為了擴展到更大的網(wǎng)表而犧牲了全局解決方案的質(zhì)量，而一個糟糕的早期分區(qū)可能會導(dǎo)致一個無法解決的最終結(jié)果；

隨機/爬山方法具有較低的設(shè)計收斂率，不能擴展到具有數(shù)百萬乃至數(shù)十億節(jié)點的現(xiàn)代芯片網(wǎng)表；

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

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

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

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

但培訓(xùn)跨芯片的 AI 布局策略非常具有挑戰(zhàn)性，因為它需要學(xué)習(xí)優(yōu)化所有可能的芯片網(wǎng)表到所有可能的芯片畫布上的布局。

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

為了實現(xiàn)泛化，研究人員重點研究了芯片的可轉(zhuǎn)移表征學(xué)習(xí)，并將表征學(xué)習(xí)應(yīng)用于預(yù)測芯片放置質(zhì)量的有監(jiān)督任務(wù)中。通過設(shè)計一種能夠準(zhǔn)確預(yù)測各種網(wǎng)表及其位置的神經(jīng)結(jié)構(gòu)，研究人員能夠生成輸入網(wǎng)表的豐富特征嵌入，然后，使用這個結(jié)構(gòu)作為編碼器，以實現(xiàn)轉(zhuǎn)移學(xué)習(xí)。

除了對芯片布局規(guī)劃的直接影響外，基于 AI 的芯片布局設(shè)計方案概括和快速生成高質(zhì)量解決方案的能力也具有重大意義，為與芯片設(shè)計過程早期階段的協(xié)同優(yōu)化創(chuàng)造了機會。

在更抽象的層次上，研究人員表示，這種形式的位置優(yōu)化策略還可出現(xiàn)在廣泛的科學(xué)和工程應(yīng)用中，例如硬件設(shè)計、城市規(guī)劃、疫苗測試和分發(fā)以及大腦皮層布局研究等。

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

設(shè)計域的自適應(yīng)策略

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

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

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

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

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

為此，研究人員首先訓(xùn)練了一個神經(jīng)網(wǎng)絡(luò)架構(gòu)，能夠預(yù)測新網(wǎng)表放置時的回報，最終目標(biāo)是將此架構(gòu)用作整個策略的編碼層，訓(xùn)練這個架構(gòu)需要一個大的芯片位置數(shù)據(jù)集和相應(yīng)的獎勵標(biāo)簽，研究人員創(chuàng)建了一個包含 10000 個芯片放置的數(shù)據(jù)集。

另外，為了準(zhǔn)確預(yù)測獎賞標(biāo)簽并推廣到未知數(shù)據(jù)，研究人員又開發(fā)了一種神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)：Edge-GNN (基于邊緣的圖形神經(jīng)網(wǎng)絡(luò))，該網(wǎng)絡(luò)的作用是嵌入網(wǎng)表，將節(jié)點類型和連通性的信息提取到低維向量表示中，以用于下游任務(wù)。

圖｜策略和價值網(wǎng)絡(luò)架構(gòu)（來源：該論文）

嵌入層會對有關(guān)網(wǎng)表鄰接、節(jié)點特征和要放置的當(dāng)前宏的信息進行編碼，然后，策略和值網(wǎng)絡(luò)分別輸出可用網(wǎng)格單元上的概率分布和當(dāng)前布局的預(yù)期回報估計。

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

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

域自適應(yīng)結(jié)果

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

研究人員展示了零樣本模式的結(jié)果，以及在特定設(shè)計上對預(yù)訓(xùn)練策略進行 2 小時和 12 小時微調(diào)后的結(jié)果。

圖｜從無到有的訓(xùn)練與不同時間量的微調(diào)，從表中可以看出，預(yù)先訓(xùn)練的策略網(wǎng)絡(luò)始終優(yōu)于從無到有訓(xùn)練的策略網(wǎng)絡(luò)，證明了從離線訓(xùn)練數(shù)據(jù)中學(xué)習(xí)的有效性（來源：該論文）

圖｜Ariane RISC-V CPU30 從零開始訓(xùn)練與從預(yù)先訓(xùn)練的策略網(wǎng)絡(luò)開始訓(xùn)練的收斂圖（來源：該論文）

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

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

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

圖｜訓(xùn)練前數(shù)據(jù)集大小的影響（來源：該論文）

當(dāng)研究人員對更多的芯片塊進行預(yù)訓(xùn)練時，能夠更快地為新的看不見的芯片塊生成更高質(zhì)量的放置。當(dāng)他們將訓(xùn)練集從 2 個塊增加到 5 個塊，最后增加到 20 個塊時，策略網(wǎng)絡(luò)在微調(diào)后都會生成更好的放置位置。

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

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

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

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

結(jié)果證明，利用 AI 進行芯片設(shè)計布局是有效的，不僅能生成高質(zhì)量的模塊放置方案，也能滿足實際設(shè)計標(biāo)準(zhǔn)要求，在面積、功率和線長方面優(yōu)于或不亞于人類專家手動放置的效果。

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

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

參考鏈接：

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

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