芯片設計效率提升2.5倍,中科大華為諾亞聯合,用GNN+蒙特卡洛樹搜索優化電路設計 | ICLR2025
芯片設計是現代科技的核心,邏輯優化(Logic Optimization, LO)作為芯片設計流程中的關鍵環節,其效率直接影響著芯片設計的整體性能。
然而,傳統邏輯優化算子由于存在大量無效和冗余的變換,導致優化過程耗時較長,成為制約芯片設計效率的主要瓶頸。
為解決這一挑戰,中科大王杰教授團隊(MIRALab)和華為諾亞方舟實驗室(HuaweiNoah’sArkLab)聯合提出了基于神經符號函數挖掘的高效邏輯優化方法,顯著提升傳統關鍵邏輯優化算子運行效率最高達2.5倍。
論文已被ICLR 2025接收。
研究團隊提出了一種創新的數據驅動的電路神經符號學習框架——CMO。
通俗解釋,研究團隊開發了一種聰明又高效的AI算法框架。這個系統通過“看圖識路”(圖神經網絡)加上“會下棋的策略”(蒙特卡洛樹搜索),能夠自動學會如何更快、更智能地“剪枝”電路邏輯——就像幫電路減肥,讓它運行得更快但不丟性能。
在實際測試中,CMO能讓關鍵算法運行效率提升最多2.5倍,也就是說,原來跑10分鐘的任務,現在4分鐘就搞定了。這個技術已經被集成進華為自研的EMU邏輯綜合工具中,有力地支持了EDA工具全鏈條國產化任務。
引言
芯片設計自動化(EDA)被稱為“芯片之母”,是半導體行業的關鍵基石。邏輯優化(LO)是前端設計流程中最重要的EDA工具之一,其核心任務是通過功能等效的轉換來優化電路,減少電路的規模和深度,從而提升芯片的質量。
LO任務是一個NP-hard問題。為了解決LO問題,現有的啟發式算子(如Mfs2[1]、Resub[2]、Rewrite[3]、Refactor[4]等)通過遍歷電路圖節點并進行局部轉換而實現邏輯優化。但由于現有算子存在大量無效和冗余的轉換,導致優化過程非常耗時,嚴重限制了芯片設計的效率。為了提高LO的效率,先前的研究提出使用打分函數來預測并剪枝無效的節點轉換。
現有的打分函數大致分為兩類,第一類是人工設計的啟發式方案[5],這些方法雖然具有可解釋性,但設計過程復雜,且泛化性能較差,難以保證算子優化性能。第二類是基于圖神經網絡(GNN)的深度學習方案[6],盡管GNN在LO任務中表現出色,但且推理嚴重依賴于GPU,難以在純CPU的工業環境中部署。此外,GNN的“黑箱”特性也引發了對其可靠性的擔憂。因此,如何找到兼具推理效率、可解釋性和泛化性能的打分函數是邏輯優化領域亟待解決的難題。
為了解決上述難題,研究團隊提出了首個數據驅動的電路神經符號學習框架(Circuit symbolic learning framework,CMO),該框架采用教師-學生范式,利用泛化性強的GNN模型作為教師,并指導作為學生的基于蒙特卡洛樹搜索的符號學習方案,從而有效生成兼具泛化能力與輕量化的符號打分函數。
在三個具有挑戰性的電路基準測試中,離線實驗結果表明,CMO學習到的可解釋符號函數在推理效率和泛化能力方面均顯著優于此前基于GPU的最先進方法以及人工設計的方案。此外,在線實驗進一步驗證了CMO的實際應用價值:CMO能夠在保持關鍵算子優化性能的同時提升其運行效率最高達2.5倍。該方案為芯片設計工具的高效化提供了新的解決方案,目前已成功應用于華為自研EMU邏輯綜合工具中。
背景與問題介紹
1.邏輯優化(Logic Optimization,LO)
邏輯優化是電子設計自動化(EDA)工具中的關鍵模塊,旨在通過優化由有向無環圖表示的電路圖(即減少電路圖的面積與深度),提升芯片的性能、功耗和面積(PPA)。邏輯優化通常分為兩個階段:
- 映射前優化(Pre-mapping Optimization):在電路映射到技術庫之前,使用啟發式算法(如Rewrite、 Resub、 Refactor等)對電路進行優化。
- 映射后優化(Post-mapping Optimization):在電路映射到技術庫(如標準單元網表或K輸入查找表)后,進一步使用啟發式算法(如Mfs2)對電路進行優化。
邏輯優化的核心任務是通過功能等效的轉換減少電路的規模和深度,從而提升芯片的質量。然而,邏輯優化是一個NP難問題,現有的啟發式算法雖然有效,但由于存在大量無效和冗余的轉換,導致優化過程非常耗時。因此,如何提高邏輯優化的效率成為芯片設計中的一個關鍵挑戰。
2.基于節點剪枝的高效邏輯優化框架
為了提升邏輯優化的效率,研究者提出了如圖1所示的預測與剪枝框架(Prediction and Prune Framework),該框架通過引入打分函數來預測并剪枝無效的節點轉換,從而減少不必要的計算開銷。具體來說:
- 節點級轉換(Node-level Transformations):邏輯優化啟發式算子(如Mfs2)通常會對電路中的每個節點依次應用轉換。然而,許多轉換是無效的,即它們不會對電路的優化結果產生實質性影響。
- 打分函數的作用:打分函數用于評估每個節點的轉換是否有效。通過預測并剪枝無效的節點轉換,可以顯著減少啟發式算法的計算量,從而提升優化效率。
△圖1.邏輯優化中的預測與剪枝框架
然而,現有的評分函數在以下幾個方面存在局限性:
1.推理效率:基于深度學習的評分函數(如GNN)雖然預測準確率高,但其推理時間長,在純CPU環境的大規模工業電路上推理時間最高可占算子運行時間的30%,因此難以滿足工業需求。2.可解釋性:深度學習模型(如GNN)的“黑箱”特性使得其在實際應用中的可靠性受到質疑。3.泛化性能:人工設計的評分函數雖然具有可解釋性,但其泛化性能較差,難以適應不同電路的特性。為了解決這些問題,本文提出了CMO框架,通過學習輕量級、可解釋且泛化能力強的符號函數來優化邏輯綜合算子。
方法
為了挖掘輕量化符號打分函數,首先提出了數據驅動的神經符號學習框架(Circuit Symbolic Learning framework, CMO),在CMO中的核心技術貢獻是圖增強的蒙特卡洛樹搜索方案(Graph Enhanced Symbolic Discovery Framework, GESD),通過學習圖神經網絡內蘊的知識以提升符號函數的泛化能力。
1.數據驅動的神經符號學習框架-CMO
如圖2所示,CMO描述了該研究整個符號函數學習與實際部署的pipeline。
△數據驅動的神經符號學習框架CMO
數據收集
通過應用邏輯優化啟發式算法(如Mfs2)對電路進行優化,生成數據集。對于電路中的每個節點,生成一個數據對,其中
是節點特征,
是標簽(如果節點轉換有效則標簽為1,否則為0)。
對于一個給定的電路圖,收集到的數據為,目標是從中學習輕量且可解釋的符號函數。
結構-語義特征分解
對于節點特征的設計,研究團隊參考了[5]的設計,將節點特征設計為了包含圖結構與圖語義信息的高維特征。其中,結構特征包含電路的拓撲信息(如節點的層級、扇入/扇出數等),而語義特征包含功能信息(如節點的真值表)。然而,高維特征會導致搜索空間急劇增加,為了解決這一問題,提出了結構-語義特征分解策略。團隊觀察到結構特征是連續的,適合數學符號回歸;而語義特征是離散的,適合布爾符號學習。因此,將特征分離并分別使用不同的符號回歸方案學習,顯著減少了符號搜索空間,并從結構與語義兩個維度集成信息,有利于模型泛化性能的提升。
神經符號函數學習
- 結構函數學習:對于連續的結構特征
,該函數將結構特征映射為連續值。
- 語義函數學習:對于離散的語義特征
,該函數將語義特征映射為離散值。
特征信息融合
在測試階段,將訓練得到的結構函數與語義函數同時作為部署模型,并將結構函數和語義函數的輸出融合,得到節點的最終分數:
其中是一個權重參數,用于平衡兩種特征的重要性。
2.圖增強的蒙特卡洛樹符號搜索方案-GESD
在CMO中,如何從給定的電路數據中發現具有強泛化能力的符號函數是一個關鍵問題。為此,提出了首個圖增強的蒙特卡洛樹符號搜索方案——GESD(如圖3所示)。
該方案通過圖神經網絡(GNN)指導蒙特卡洛樹的生成,巧妙結合了圖神經網絡的高泛化能力與符號函數的輕量化優勢,從而顯著提升了符號函數的泛化性能。
△圖3. 圖增強的蒙特卡洛樹符號搜索方案GESD
符號樹生成
- 符號操作符:在生成符號函數之前,需要定義搜索中使用的符號。用于生成結構表達式樹的數學運算符包括:{+, ?, ×, ÷, log, exp, sin, cos}。團隊沒有使用占位符來生成常數,因為引入內部常熟優化循環通常會導致 更高的訓練成本。此外,團隊發現像exp、sin、cos等復雜操作符能夠有效提升符號函數的預測性能。對于生成語義表達式樹的布爾運算符包括:{與,或,非},通過與或非的組合能夠擬合任意一個布爾表達式。
- 蒙特卡洛樹搜索:受到蒙特卡洛樹搜索(MCTS)在有效探索大型復雜符號空間方面的優勢啟發(Sun等,2023;Xu等,2024),使用符號表達式樹來表示符號函數通過表達式樹并利用MCTS生成符號樹。對于一棵符號樹,他的內部節點是數學運算符(如加、減、乘、除、對數、指數等),葉子節點是輸入變量或常數。定義狀態(s)為當前表達式樹的先序遍歷,動作(a)為添加到狀態中的符號操作符或變量。具體來說,此研究中的MCTS包括四個步驟:選擇、擴展、模擬和反向傳播。
1.選擇:在選擇階段,MCTS代理遍歷當前表達式樹,選擇具有最大UCT值的動作。為了確保生成表達式的合法性,在當前狀態下,MCTS代理會屏蔽掉當前非終端節點的無效動作,并在此基礎上選擇一個有效的動作作為。2.擴展:當選擇階段達到一個可擴展節點——即它的子節點并非全部已被訪問——MCTS代理會通過隨機選擇一個未訪問的有效子節點來擴展該節點。3.模擬:在當前狀態和擴展節點的基礎上,通過隨機選擇下一個節點進行仿真,直到表達式樹完成。具體來說,進行10次仿真,并返回最大仿真獎勵,而不是傳統MCTS算法中的平均獎勵,以找到唯一的最優符號解,這是一種與傳統MCTS算法不同的貪婪搜索啟發式方法。4.反向傳播:仿真結束后,更新路徑中從當前節點到根節點的最大獎勵Q和訪問次數N。該搜索算法會不斷重復上述步驟,直到滿足停止準則。
圖增強符號函數
- 教師-學生框架:
引入圖神經網絡(GNN)作為教師模型,通過蒸餾GNN的“暗知識”來增強符號函數的泛化能力。具體來說,首先在訓練數據集上訓練一個GNN,該GNN能夠有效捕捉電路中的復雜結構和語義信息,從而解決由于電路領域分布大幅度變化而導致的泛化能力差的問題。然后,利用GNN的預測輸出和真實標簽來指導符號函數的學習過程。
這種教師-學生框架的核心思想是通過GNN的高泛化能力來引導符號函數的生成,從而彌補傳統符號方法在泛化性能上的不足。
- 獎勵函數設計:
在MCTS的模擬階段,使用以下獎勵函數評估符號函數:
其中是基于標簽的損失,用于確保符號函數的預測結果與真實標簽一致;
是基于GNN輸出的損失,用于將GNN的泛化能力遷移到符號函數中。通過調整權重
可以平衡標簽信息和教師知識的重要性。此外,
是一個懲罰因子,用于控制符號函數的復雜度,確保生成的符號函數既簡潔又高效。
- 圖蒸餾
通過最小化符號函數輸出與GNN輸出之間的均方誤差(MSE),將GNN的泛化能力遷移到符號函數中。與以往使用KL散度的方法不同,MSE能夠直接學習GNN輸出的泛化信息,從而實現與GNN相當的泛化性能。具體來說,研究團隊發現電路特征與GNN輸出之間存在簡單的非線性映射關系,這使得符號函數能夠在不犧牲性能的情況下保持輕量化。
此外,針對電路數據中正負樣本嚴重不平衡的問題,采用焦點損失(Focal Loss)作為學生模型的損失函數,進一步提升符號函數的學習效果。焦點損失通過調整難易樣本的權重,有效緩解了樣本不平衡帶來的負面影響,從而提高了符號函數在稀疏數據上的表現。
實驗結果
在實驗部分,在兩個廣泛使用的開源電路數據集(EPFL和IWLS)以及一個工業電路數據集上進行了全面測試。實驗結果表明,該方法在多個方面表現出色:
1.效率提升:CMO框架顯著提升了傳統關鍵邏輯優化算子(如Mfs2)的運行效率,最高可達2.5倍加速。例如,在超大規模電路Sixteen上,CMO-Mfs2將運行時間從78,784秒減少到32,001秒,提升了約59.4%。
2.優化質量提升:通過在相同時間內多次運行CMO驅動的新型算子(如2CMO-Mfs2),進一步提升了電路的優化質量(QoR)。具體來說,電路的規模和深度得到了顯著改善,其中電路深度的最大優化幅度達到30.23%。例如,在Hyp電路上,2CMO-Mfs2將電路深度從8,259層減少到5,762層,顯著降低了電路的延遲。
這些實驗結果充分證明了CMO框架在提升邏輯優化效率和質量方面的雙重優勢,為芯片設計中的邏輯優化任務提供了強有力的支持。
△表1.實驗結果表明該方法不僅能夠提升算子效率,同時還能夠提升算子優化質量
本論文作者白寅岐是中國科學技術大學2024級碩士生,師從王杰教授,主要研究方向為人工智能驅動的芯片設計、圖機器學習、大模型等。他曾在人工智能頂級會議ICML、Neurips等會議上發表論文兩篇,本科期間曾獲首批國家青年學生基礎研究項目資助(全國108人)。
論文地址:https://openreview.net/forum?id=EG9nDN3eGB
代碼地址:https://gitee.com/yinqi-bai/cmo