計關(guān)鍵:如何把8TOPS跑出高利用率)
最近幾年做AI芯片的項目越來越多但很多團隊做出來的片子標(biāo)桿測試很好看一上真實模型就掉鏈子。問題往往不在某個硬件模塊本身而是整條鏈路上軟硬件到底怎么切、怎么協(xié)同。我參與過一個邊緣側(cè)推理芯片的設(shè)計規(guī)模不大目標(biāo)8TOPS算力跑主流檢測模型端到端延遲壓在30毫秒以內(nèi)。今天就用這個例子把AI芯片軟硬件設(shè)計的關(guān)鍵環(huán)節(jié)拆開聊聊從算力指標(biāo)怎么定到算子怎么映射再到存儲墻和互聯(lián)設(shè)計最后是軟件棧怎么配合。內(nèi)容不涉及具體產(chǎn)品但流程和思路是可以復(fù)用的。先說一個基本判斷AI芯片設(shè)計的核心不是“堆算力”而是“把用戶程序跑出高利用率”。8TOPS的峰值算力只是一個包裝值真正決定產(chǎn)品好壞的是有效算力——也就是跑實際網(wǎng)絡(luò)時能穩(wěn)定達(dá)到的利用率。而這個利用率幾乎全部是由軟硬件切分策略決定的。1. 算力指標(biāo)不是越高越好關(guān)鍵是能不能用滿做AI芯片第一個被問的就是“峰值算力多少”。8TOPS聽起來很明確但這個數(shù)字背后有一整套假設(shè)乘加位寬是INT8還是FP16算不算稀疏加速是單核跑滿還是多核并行這些口徑不同數(shù)字可以差出好幾倍。我在項目里定指標(biāo)的方式是反著來的先明確目標(biāo)模型和延遲再倒推有效算力。比如要在30毫秒內(nèi)跑完一個單幀1080p檢測主算子大約6.2G次乘加運算。如果不考慮任何開銷光算MAC的等效算力需求是30ms處理6.2GMAC折合每秒約206.7 GMAC/s也就是約413 GOPS乘加算2次浮點或定點運算??紤]到實際利用率不可能到100%按設(shè)計目標(biāo)70%利用率反推峰值算力至少要590 GOPS。而如果按INT8實現(xiàn)就是大約8TOPS——我們后來定的“8TOPS”這個宣傳數(shù)字其實是這么算出來的。為什么強調(diào)口徑因為很多人拿到一個標(biāo)稱“8TOPS”的芯片以為真能跑出8T的吞吐結(jié)果實測下來只有2到3T然后罵硬件不給力。這不全是硬件的鍋——通常是你跑的算子形狀和硬件排布不匹配或者存儲帶寬跟不上或者DMA調(diào)度和計算沒重疊。硬件只是提供了一個潛力上限能不能兌現(xiàn)靠的是整個軟硬件系統(tǒng)。還有一個容易忽略的點INT8和FP16的算力標(biāo)稱完全是兩回事。如果芯片聲稱8TOPS一定要問清是INT8還是FP16。FP16的8TOPS做成INT8通常會翻倍取決于架構(gòu)是否支持而實際部署中很多模型還需要混合精度不能只看一個數(shù)字。2. 從卷積到MAC陣列算子映射是需求翻譯機AI芯片硬件設(shè)計的邏輯起點不是“我要做一個8核處理器”而是“我的算子要在什么形狀下跑”。以卷積為例一個標(biāo)準(zhǔn)的3×3卷積滑動窗口在空間和通道維度不斷復(fù)用數(shù)據(jù)。如果你拿CPU或GPU的思路去設(shè)計很容易陷入“每個線程算一個輸出點”的慣性。但AI芯片不是這樣——它更適合做的是“把一批輸出通道的計算打包給一個MAC陣列”。我們當(dāng)時選擇的映射策略是輸出通道做并行切分輸入通道做串聯(lián)累加。一個16×16的MAC陣列硬件上就是一列算16個不同輸出通道一行做16個輸入通道的乘加累加。這樣適配起來非常直接。然后問題變成輸入圖像切成多少個tile、每次塞多少個輸入通道可以既保持陣列滿負(fù)荷又不會讓SRAM爆炸。這里有個實操經(jīng)驗與其寫一堆復(fù)雜的調(diào)度算法不如先做一個簡單的“純卷積峰值測試”——把所有輸入輸出固定在片上不經(jīng)過DMA看MAC陣列到底能跑多滿。如果這個測試也只能跑到60%那硬件設(shè)計大概率有問題如果能跑滿說明問題出在數(shù)據(jù)搬運和同步上。這個測法幾乎不花時間卻能幫你快速分清瓶頸是計算還是存儲。我們還研究過一個細(xì)節(jié)卷積核的形狀是否固定。很多AI芯片為了省事只能跑3×3其他形狀全部拆成1×1組合。這在小網(wǎng)絡(luò)里還勉強能用但遇到大模型、大感受野時拆解帶來的數(shù)據(jù)重排開銷會吃掉很多性能。所以我們的陣列設(shè)計支持通用矩陣乘的抽象不管是卷積還是全連接還是Transformer里的線性投影都統(tǒng)一映射到矩陣乘。這套方案在工程實現(xiàn)上稍微復(fù)雜但收益明顯——之后適配新模型時不用改硬件。3. 算力越強存儲墻越要命數(shù)據(jù)搬運才是真正的瓶頸我見過不少團隊把MAC陣列做大結(jié)果性能反而沒漲多少。深挖之后發(fā)現(xiàn)問題幾乎都出在存儲墻上計算一次只要幾十個周期但取一次數(shù)據(jù)要幾百個周期。在IC設(shè)計里這叫“memory-bound”幾乎每個AI芯片都逃不開。片上存儲的設(shè)計思路我總結(jié)為三層遞進(jìn)第一層是寄存器/累加器這是離MAC最近的數(shù)據(jù)歸屬速度最快但容量極小。第二層是SRAM也就是常說的片上緩存用來放輸入tile、權(quán)重tile和輸出tile容量在幾百KB到幾MB之間。第三層是片外DRAM或者HBM這個容量大但延遲和帶寬都不便宜。關(guān)鍵設(shè)計動作是“數(shù)據(jù)復(fù)用”。我用一個簡單的估算來說明假設(shè)SRAM帶寬是256字節(jié)/周期主頻1GHz那每秒能從SRAM讀出256GB數(shù)據(jù)。如果8TOPS的MAC陣列跑滿每周期要消費64字節(jié)的輸入和64字節(jié)的權(quán)重才能喂飽。這組數(shù)字雖然粗糙但反應(yīng)了一個關(guān)系MAC的算力增長是線性的存儲帶寬的工程代價卻是超線性的。你多塞一組MAC就要多鋪一堆SRAM讀寫端口最后片上布線都成了問題。實際項目里我們采用了“tiling double buffering”的組合策略。tiling就是把大特征圖切成小塊盡量讓一個小圖在SRAM里完成所有計算再搬走double buffering就是讓DMA搬運下一塊數(shù)據(jù)的同時MAC陣列計算當(dāng)前這塊。這兩個手段疊加大約能把搬運時間隱藏掉60%以上。剩下的不能隱藏的就是你MAC利用率的天花板。還一個容易被忽視的地方輸出數(shù)據(jù)的寫回。大多數(shù)人只關(guān)心輸入怎么搬進(jìn)來忽略輸出怎么搬出去。如果網(wǎng)絡(luò)很寬比如輸出通道很多寫回帶寬不夠照樣阻塞流水線。我們吃過一次虧某個卷積層的輸出tile特別大結(jié)果DMA讀和寫同時搶帶寬整體性能掉了一半。后來在調(diào)度里加了讀寫通道的優(yōu)先級控制才算解決。3.1 為什么說“算得不夠快”往往是“搬得不夠快”我用一個生活化的比方來解釋這個道理。假設(shè)你是一個廚師爐子MAC陣列一分鐘能炒10道菜但傳菜員DMA一分鐘只能端過來5道菜的原料那你再怎么提高廚藝實際出菜速度就是5道。AI芯片的很多性能問題不是“計算不夠快”而是“原料供應(yīng)跟不上”。這個“原料供應(yīng)”包含兩部分一是從片外DRAM到SRAM的搬入二是SRAM的訪問帶寬。前者是大動脈后者是毛細(xì)血管。兩者任何一個塞住MAC都得空轉(zhuǎn)。設(shè)計階段一定要把“數(shù)據(jù)流圖”畫出來標(biāo)注每個算子的輸入大小、輸出大小、復(fù)用次數(shù)然后逐級核對帶寬預(yù)算。寧可早一點發(fā)現(xiàn)瓶頸并調(diào)整拍頻也別等流片回來再抓瞎。4. 片上互聯(lián)不只是布線它決定了多核能不能一起干活做到這一步很多項目會卡在“多核協(xié)同”上。芯片內(nèi)部通常有多個計算核每個核都有自己的MAC陣列和SRAM核與核之間要交換數(shù)據(jù)——這就是片上網(wǎng)絡(luò)NoC的職責(zé)。我看到很多小團隊在這一步犯難要么懶得設(shè)計直接用一根共享總線便宜省事但多核同時訪問時帶寬互相打架要么設(shè)計得太重用了復(fù)雜的mesh網(wǎng)絡(luò)但實際負(fù)載根本用不上那么多路徑。AI芯片的互聯(lián)設(shè)計不應(yīng)該追求通用系統(tǒng)的靈活性而要貼合自己的并發(fā)模型。我們的設(shè)計目標(biāo)很樸素讓每個核的大部分?jǐn)?shù)據(jù)在本地SRAM里自給自足只有必須跨核的數(shù)據(jù)才走NoC。所以NoC的壓力被大幅降低帶寬需求實際上很小。為了讓這個目標(biāo)能落地算子切分時就要考慮“每個核的輸出能否直接用在后處理里”盡量減小跨核通信的頻次。還有一個細(xì)節(jié)是同步開銷。多核并行最怕“等你等它”一個核算完要等另一個核的數(shù)據(jù)才能繼續(xù)。如果同步周期太長并行度再高也白搭。我們在硬件上加了barrier同步單元讓核間的等待控制在幾十個周期內(nèi)這樣實際并行效率能到85%以上。開發(fā)時也請軟件同事提前介入讓他們清楚知道核間同步的粒度寫算子時能主動避免頻繁的全局同步。再提一句布線布局。AI芯片面積有限SRAM占了一大塊走線的長度和擁塞程度直接決定最高主頻。如果為了好看把NoC設(shè)計得極其規(guī)整但導(dǎo)致關(guān)鍵路徑過長主頻上不去那還不如降復(fù)雜度保頻率。硬件設(shè)計始終是平衡的藝術(shù)沒有放之四海而皆準(zhǔn)的方案。5. 軟件棧不是插線板編譯器、算子庫和運行時必須和硬件一起長如果只把軟件當(dāng)成“事后的解釋器”這個項目十有八九會延期。AI芯片的軟硬件設(shè)計“軟”字占一半而且必須是同步設(shè)計——硬件還在寫RTL的時候軟件團隊就應(yīng)該開始搭編譯器框架和算子庫了。算子庫的角色是“讓常見算子直接跑得不錯”。我們?yōu)槊總€常用算子手寫了優(yōu)化版本比如3×3卷積、1×1卷積、深度卷積、矩陣乘、歸一化、池化等。手寫算子的好處是你能精確控制數(shù)據(jù)排布和循環(huán)變換。但手寫也有代價——新算子層出不窮不可能全都人工優(yōu)化。所以算子庫的定位應(yīng)該是一個“含有一組高質(zhì)量模板”的集合而不是一個封閉的靜態(tài)的庫。編譯器的任務(wù)則是“把新算子壓進(jìn)模板的審美體系里”。它做兩件事圖優(yōu)化和指令映射。圖優(yōu)化負(fù)責(zé)算子的融合、改寫和內(nèi)存復(fù)用指令映射負(fù)責(zé)把算子的循環(huán)結(jié)構(gòu)映射到MAC陣列的指令序列。這套流程有點像“翻譯官”硬件是個只懂幾種固定句式的老外編譯器要把用戶的任意算法轉(zhuǎn)換成老外能聽懂的短句。這里有一個很關(guān)鍵的經(jīng)驗編譯器不要只依賴自動調(diào)優(yōu)因為自動搜索空間的邊界往往就是硬件架構(gòu)的邊界搜索不到就是搜索不到。我的做法是手動調(diào)優(yōu)做主路徑自動搜索只負(fù)責(zé)在邊界附近小幅擾動。舉個例子固定矩陣乘的映射方式為“輸出通道并行、輸入通道累加”然后讓編譯器在tile大小、緩沖位置、讀寫順序上自動搜索。這樣既能保證線上穩(wěn)定又能榨出最后幾個百分點的性能。運行時runtime也不能忽略。它調(diào)度DMA、管理多核分配、控制功耗和溫度。如果你把硬件設(shè)計成“必須由運行時精確配置才能工作”的狀態(tài)那運行時復(fù)雜度和硬件耦合度會急劇上升反而容易出錯。我傾向于硬件的默認(rèn)狀態(tài)就能工作運行時只做優(yōu)化性決策這樣軟件出了bug也不會導(dǎo)致硬件完全不可用。5.1 一個“算子先做還是硬件先定”的爭議怎么破項目里經(jīng)常會有這種爭論硬件團隊說“你先把算子需求發(fā)我”軟件團隊說“你先把寄存器模型發(fā)我”兩邊互相等。拆解這個問題關(guān)鍵是分清楚哪些需求是長期不變的“錨點”哪些是可以妥協(xié)的“變量”。錨點是那些在絕大多數(shù)模型中都有、且變化不大的計算模式卷積、矩陣乘、逐元素操作、池化、歸一化。這些必須在硬件設(shè)計的第一周就定下來連MAC陣列的尺寸、累加器的位寬都要圍繞它們來設(shè)計。變量則是具體網(wǎng)絡(luò)里的特殊層比如動態(tài)形狀的分支、非常規(guī)的激活函數(shù)。這些可以走編譯器標(biāo)量處理或者CPU核去兜底不需要立刻定死。我們的做法是用一版“虛擬指令集時鐘級行為模型”做軟硬件接口。硬件團隊按這個接口做RTL軟件團隊按這個接口寫模擬器上的算子原型。兩邊不互相等但接口定得很細(xì)致連每次DMA搬運的配置字段都提前約束好。這樣到仿真階段幾乎不會出現(xiàn)“底層怎么不按約定實現(xiàn)”的爭吵。6. 實測階段最重要的三件事先測邊界再測融合最后測長尾設(shè)計和驗證做完終于到了FPGA原型或流片回來。這個時候最容易犯的錯誤是一上來就整模型跑benchmark。我建議反過來先做三個層層遞進(jìn)的實測。第一測邊界每個算子單獨跑用峰值數(shù)據(jù)把MAC陣列打滿測真實的最大算力。如果這個數(shù)字達(dá)不到設(shè)計值先別怪軟件回頭查流水線、查memory端口沖突、查陣列的復(fù)位邏輯。邊界測穩(wěn)定了才有資格談后續(xù)。第二測融合把兩個相鄰算子連接起來比如卷積歸一化測一下數(shù)據(jù)是否能先做個融合減少一次tile搬運。這里最容易暴露DMA調(diào)度和計算流的疊合能力。我在項目里曾把融合后的算子性能提升了30%靠的就是把歸一化的縮放系數(shù)直接融進(jìn)卷積的偏置項省掉了一次整層的讀寫。第三測長尾跑真實的完整模型尤其是那種分支多、shape不規(guī)則的模型。測試時一定要記錄每一層實際消耗的周期和利用MAC周期的比例而不是只看整個模型跑完的毫秒數(shù)。很多性能黑洞出現(xiàn)在看似不起眼的小層上比如reshape、split、concat這些“操作型算子”它們在mac陣列上可能根本不做計算但數(shù)據(jù)搬移的時間卻非??捎^。這組實測做完基本能定位出整個系統(tǒng)的真實瓶頸再決定是優(yōu)化編譯器、調(diào)整runtime還是修改硬件。6.1 容易忽略的“memory layout”坑最后分享一個我們踩了很久的坑希望后來者繞開數(shù)據(jù)排布方式memory layout。同一個特征圖是按“通道-高-寬”存還是按“高-寬-通道”存對某些算子的親和度完全不同。卷積喜歡連續(xù)通道逐元素操作喜歡連續(xù)元素。如果你的算子庫內(nèi)部數(shù)據(jù)排布不統(tǒng)一那每次轉(zhuǎn)換都是一次隱形的數(shù)據(jù)搬運。我們前期因為圖省事讓不同算子各自決定layout結(jié)果引入了大量轉(zhuǎn)換層。后來統(tǒng)一成模型推理時全部用“通道優(yōu)先塊化”的方式存儲權(quán)重在編譯期重排成硬件友好的格式而不是運行期轉(zhuǎn)換。這個決定讓整體性能提升了接近兩成也大大簡化了runtime的邏輯。AI芯片的軟硬件設(shè)計不是一個“先硬件后軟件”的串行流程更像一套骨架和肌肉同步發(fā)育的過程。算力指標(biāo)要反推算子映射要精算存儲和互聯(lián)要圍繞數(shù)據(jù)流來塑形軟件棧要提前進(jìn)場而實測時還得有耐心一輪輪去摳細(xì)節(jié)。這個領(lǐng)域沒有捷徑但每一個環(huán)節(jié)想清楚了做出來的東西會明顯比對手更“耐跑”。我自己的最大體會是把“效率”當(dāng)作設(shè)計的第一性指標(biāo)而不是“規(guī)格”。規(guī)格是給市場聽的效率才是你和客戶真正長期相處的東西。