同設(shè)計:算子映射、指令集與編譯器適配全解析)
AI芯片這行聊架構(gòu)的文章很多聊編譯器和指令集的文章也不少但能把“軟硬件設(shè)計”串成一條線講清楚的確實不多。這個系列我寫到了第6篇前幾篇我把脈動陣列、片上網(wǎng)絡(luò)、存儲層次這些硬件架構(gòu)翻來覆去講過一遍。結(jié)果不少做軟件的朋友留言說你們硬件選型倒是講得熱鬧可這些東西到底怎么跟編譯器、算子庫配合為什么同一塊芯片有人能跑到理論峰值的80%有人拼了命只能用到三成這篇文章我就把軟硬件協(xié)同設(shè)計里最容易被忽略、但決定成敗的幾個環(huán)節(jié)——算子映射、指令集裁剪、編譯器適配、流片前聯(lián)合驗證——挑出來逐個拆開。適合正在做AI芯片工具鏈、或者準(zhǔn)備從算法轉(zhuǎn)芯片方向的人讀如果你剛?cè)腴T也不用慌我會先講清楚“為什么必須要軟硬件協(xié)同”再把每一環(huán)背后的邏輯說透。1. 軟硬件劃分AI芯片設(shè)計的第一個岔路口1.1 為什么說架構(gòu)決定軟件的天花板很多人以為軟硬件協(xié)同是從寫代碼那一刻開始的其實不是真正的分岔路口在芯片定義階段就出現(xiàn)了。你要先想清楚哪些計算被固化到硬件里哪些留給軟件去靈活處理。這個劃分決定了后面的軟件團(tuán)隊在幾層樓里跳舞也決定了你的芯片兩年后是“一把梭”還是“生銹的老古董”。拿卷積舉例。卷積運(yùn)算在圖像和AI模型里占了九成以上的計算量錘死在硬件里很劃算。但你要是把所有的卷積細(xì)節(jié)、不同的padding方式、不同尺寸的kernel全部用純硬件邏輯實現(xiàn)那芯片面積會爆掉功耗也會失控。而且算法迭代特別快今天流行卷積明天可能流行注意力機(jī)制或者新的算子。硬件一旦固死了就只能看著新模型干瞪眼。所以業(yè)界的普遍做法是把最底層的乘累加陣列、激活函數(shù)、池化這類“原子操作”做成硬件單元而把“如何組合這些原子操作”留給了指令流。這就是軟硬件劃分的核心硬件負(fù)責(zé)確定性高、重復(fù)性高、能耗敏感的部分軟件負(fù)責(zé)策略性、組合性和未來可擴(kuò)展的部分。我參與過某AI推理加速芯片項目當(dāng)時團(tuán)隊就為“是否把softmax也做成硬件單元”吵了一星期。做硬件的人說softmax里那個exp和除法太費電必須硬化做軟件的人說模型里經(jīng)常有自定義變體硬死了不好改。最后方案是把exp查表做成了硬件小單元但softmax的整體流程交給軟件算子庫去拼裝。這樣既保住了大部分功耗收益又留了靈活性。1.2 劃分時的取舍原則做軟硬件劃分時我建議死守三條原則。第一看“調(diào)用頻率”。頻率極高的算子例如矩陣乘、卷積、向量加法果斷傾向硬化一年用不上幾次的算子就算再復(fù)雜也別為它浪費芯片面積。第二看“數(shù)據(jù)依賴復(fù)雜度”。像卷積這種數(shù)據(jù)復(fù)用規(guī)律非常清晰的計算硬化后可以用脈動陣列把數(shù)據(jù)搬移成本壓到極低而像動態(tài)維度、動態(tài)shape這類運(yùn)行時才能確定的東西硬到硬件里只會讓你想砸鍵盤。第三看“生態(tài)兼容窗口”。你做的芯片總要跑主流框架吧哪怕今天某個算子不常用但如果它長在主流模型的骨干路徑上那硬件還得多留一步軟實現(xiàn)的后路。這里有個容易踩的坑過度適配。我見過某芯片把三年前的某小眾模型算子固化得非常漂亮性能好看得很但過兩年那個模型已經(jīng)沒人用了芯片反而因為那部分硬件占用沒法給新算子騰地方。軟硬件劃分的本質(zhì)是風(fēng)險管理不是炫技。2. 算子映射與指令集設(shè)計讓硬件聽懂算法2.1 從Tensor到指令算子映射的路徑假設(shè)你已經(jīng)從算法里抽出來了一個矩陣乘算子下一步就是要把這個算子翻譯成芯片能執(zhí)行的指令序列。這段路大多數(shù)人叫“算子下發(fā)”但里面門道很深。首先你要把矩陣乘這個邏輯操作拆成循環(huán)嵌套遍歷輸出矩陣的行、列然后做內(nèi)積。針對每個芯片的計算陣列規(guī)模循環(huán)要切塊tiling。比如你的片上陣列是16x16的乘法單元但你算的矩陣是128x128那就得把大矩陣切成一堆16x16的小塊再一塊塊喂進(jìn)去。這個切塊的大小直接影響兩個東西一是數(shù)據(jù)復(fù)用小塊從片外搬上來越少越好二是流水線銜接上一塊還在算下一塊的數(shù)據(jù)最好已經(jīng)在片上等著了。很多芯片利用率低不是因為計算陣列不行而是因為切塊策略一塌糊涂算一會兒等一會兒等于花錢養(yǎng)了一堆閑著的乘法器。我舉個例子某芯片的矩陣單元是32x32的MAC陣列理論算力挺高。結(jié)果跑一個維度為512的全連接層時軟件團(tuán)隊直接把整個矩陣按一個小塊硬拆片外拷了四次數(shù)據(jù)把時間全浪費在DDR帶寬上了。后來改成混合分塊行方向按32拆列方向按64拆讓片上數(shù)據(jù)復(fù)用率達(dá)到70%以上實測端到端延遲直接降了三分之一。2.2 指令集擴(kuò)展的取舍把算子拆分好了還得靠指令把它表達(dá)出來。早期很多AI芯片直接沿用CPU的SIMD指令后來發(fā)現(xiàn)矩陣乘用SIMD指令得拆成無數(shù)條向量乘加指令流量太大跑到飛起性能也上不去。于是大家開始做指令集擴(kuò)展。一種是VLIW風(fēng)格一條超長指令同時控制多個功能單元并行工作。好處是編譯器可以做大量調(diào)度壞處是解碼器復(fù)雜、指令字長驚人很容易把指令存儲空間吃滿。一種是張量指令風(fēng)格就是你現(xiàn)在經(jīng)常聽到的“一條指令算一個矩陣乘”。這種方法指令密度高但對硬件的數(shù)據(jù)通路要求也高通常得配合專門的脈動陣列或二維寄存器堆。我和不少同行聊過現(xiàn)在主流趨勢其實是混合模式控制密集、不規(guī)則的部分用通用指令計算密集、規(guī)則的部分用張量指令。關(guān)鍵是接口要干凈別讓軟件每次下發(fā)都要猜硬件狀態(tài)。我看過一些芯片手冊指令集設(shè)計得極其豐富但編譯器只用了其中三分之一剩下全是沒用的死功能。指令集不是越多越強(qiáng)而是找出現(xiàn)實模型里最常出現(xiàn)的計算模式把那些模式啃下來就夠了。這個取舍在AI芯片上比傳統(tǒng)處理器要激進(jìn)得多——因為AI芯片的指令集是和一個相對封閉的模型生態(tài)打交道的少而精準(zhǔn)更香。3. 編譯器與運(yùn)行時把高層語言落到底層硬件3.1 編譯器的層次與分工指令集的另一端就是編譯器。AI芯片的編譯器和傳統(tǒng)編譯器長得不太一樣它更像一個“調(diào)度器優(yōu)化器代碼生成器”的組合體。一般分三層干活前端負(fù)責(zé)把PyTorch、TensorFlow出來的計算圖接進(jìn)來做算子融合。比如把卷積后面的批量歸一化、ReLU融合進(jìn)同一個算子減少一次中間結(jié)果落回內(nèi)存。很多新手工改框架后端往往只盯著單個算子跑得快不快完全忘了算子之間的數(shù)據(jù)搬運(yùn)才是真正的消耗大戶。融合一次可能比單個算子優(yōu)化半天的收益還大。中端負(fù)責(zé)和硬件參數(shù)打交道做循環(huán)分塊、數(shù)據(jù)布局轉(zhuǎn)換、內(nèi)存分配。這一層是最容易翻車的。比如二維數(shù)據(jù)在內(nèi)存里是按行還是按列排直接決定搬運(yùn)效率和向量化效率。很多算法模型用NCHW格式但你的硬件如果對NHWC更友好編譯器就得自動插入layout轉(zhuǎn)換這個轉(zhuǎn)換本身也有成本如何均衡是門手藝活。后端就是所謂的指令調(diào)度。把已經(jīng)標(biāo)好的操作排序、綁定到具體的執(zhí)行單元、處理好流水線停頓。如果前中端做得好后端壓力會小很多反之后端再精心打補(bǔ)丁也沒用。有個觀點我想多說一句AI芯片編譯器里永遠(yuǎn)不要指望“一招吃遍天”。不同模型、不同尺寸的算子最優(yōu)調(diào)度方案可能截然不同。與其在編譯器里堆一堆規(guī)則不如建立一個靜態(tài)性能模型讓編譯器基于成本模型做決策。我在某項目里就用了個非常簡化的帶寬估算模型就把不少算子的調(diào)度決策從“拍腦袋”改成了“按數(shù)據(jù)計算”效果立竿見影。3.2 算子調(diào)度與內(nèi)存分配講到具體操作大部分性能問題都繞不開存儲層次。AI芯片基本都有寄存器堆、SRAM、DDR這幾級存儲。數(shù)據(jù)從DDR搬到SRAM要幾十上百個周期從SRAM搬到寄存器堆只要幾個周期差距懸殊。所以核心原則就一句話盡可能讓數(shù)據(jù)在離計算單元近的地方多被復(fù)用幾次。我做算子調(diào)優(yōu)時常用的手段是“雙重循環(huán)分塊”先按輸出通道分塊讓同一塊輸入數(shù)據(jù)對應(yīng)多個輸出通道的計算再按輸入通道分塊讓同一輸出區(qū)域復(fù)用多組輸入。這樣做能顯著增加SRAM里的數(shù)據(jù)命中次數(shù)而不是每算一個小塊跑去DDR搬數(shù)據(jù)。內(nèi)存分配呢最大的坑是“內(nèi)存碎片”。AI芯片的SRAM通常不大幾兆已經(jīng)算多的但模型一進(jìn)來各種張量同時并存如果你用簡單的先到先得策略很容易把SRAM拆得七零八落最后大模型根本裝不下。我見過做得好的編譯器會先用一個“模擬退火貪心”的分配器搜索一遍把生命周期相近的張量安排到重疊區(qū)域然后把剩余空間整合給最大的中間結(jié)果。這一步做扎實了某些模型占用面積能減少四成。3.3 運(yùn)行時與驅(qū)動的設(shè)計編譯出來的“可執(zhí)行文件”只是一堆指令和數(shù)據(jù)的集合真正讓它們在芯片上轉(zhuǎn)起來的是運(yùn)行時和驅(qū)動。運(yùn)行時干的事很雜內(nèi)存管理、任務(wù)隊列、同步、性能計數(shù)器的上報。我印象最深的是任務(wù)啟動成本。很多芯片每次下發(fā)一個小算子驅(qū)動就要做一堆寄存器配置、內(nèi)存映射、中斷處理。如果一個模型有幾百個小算子每秒跑幾百輪光啟動開銷就把性能吃光了。這也解釋了為什么現(xiàn)在的落地產(chǎn)品格外強(qiáng)調(diào)“大算子fusion 一次下發(fā)”把幾百個小算子揉成幾個大任務(wù)啟動開銷一下子降到可以忽略。另外別小看了性能計數(shù)器的設(shè)計。調(diào)優(yōu)的時候沒有精確的cycle數(shù)據(jù)你只能瞎猜瓶頸在哪。給運(yùn)行時加上指令計數(shù)、緩存miss計數(shù)、DRAM讀寫帶寬統(tǒng)計讓軟件團(tuán)隊能拿到多條硬件流水線的吞吐和占用率這是軟硬件協(xié)同能否持續(xù)優(yōu)化的關(guān)鍵。某芯片就是因為在硬件里埋了一組輕量級的性能監(jiān)控點之后編譯器團(tuán)隊才有了明確的優(yōu)化方向。4. 軟硬件協(xié)同驗證流片前如何提前發(fā)現(xiàn)問題4.1 性能建模與仿真的位置軟硬件設(shè)計最大的痛點是軟件開發(fā)想提早進(jìn)行但芯片還沒流片。傳統(tǒng)的做法是各干各的——硬件做RTL仿真軟件做指令模擬器兩邊完全脫節(jié)到了板子上才發(fā)現(xiàn)軟硬件接口對不上一改就是幾個月周期根本等不起。我現(xiàn)在強(qiáng)烈建議在項目初期就建立一套“性能模型指令集模擬器”的組合。性能模型不需要拿到精確到信號級的延遲但要用會計量數(shù)據(jù)在每一級存儲上的駐留時間、算子執(zhí)行單元的占用情況。它能在一小時內(nèi)驗證“某種切塊策略會不會導(dǎo)致長時間等待”給指令集設(shè)計提需求。指令集模擬器則更精細(xì)一點執(zhí)行具體的指令流輸出計算和搬運(yùn)事件的時序。兩套工具并行相當(dāng)于在設(shè)計階段就跑通了軟硬件接口。好多人跟我抱怨說性能建模工作量太大做不完。我的體會是別沉浸在精確度里先做能算“相對瓶頸”的工具。你需要回答的其實只有三個問題數(shù)據(jù)搬運(yùn)需要多少個cycle計算單元有多少cycle是空閑的到底誰在拖誰的后腿有了這三個答案大部分架構(gòu)問題就現(xiàn)形了。4.2 虛擬原型與聯(lián)合調(diào)優(yōu)等指令集和編譯器初版都有了更完整的驗收到虛擬原型Virtual Prototype階段。這一層可以用軟件模擬器把RTL的行為模型化讓完整應(yīng)用在上面跑甚至把調(diào)度器的策略放進(jìn)去觀察效果。我在某項目中做過一次實驗同樣是卷積算子在虛擬原型上換了一種存儲排列順序利用率從56%升到78%而這次調(diào)整發(fā)生在流片前半年相當(dāng)于白賺了一輪性能優(yōu)化。不過虛擬原型也有自己的問題速度慢。給你一個調(diào)試用的慢版本可以但別指望它大規(guī)模跑訓(xùn)練或推演整套手機(jī)端性能。所以在實踐中我會把驗證策略分成三檔單算子級別跑架構(gòu)級模擬器小網(wǎng)絡(luò)級別跑虛擬原型完整應(yīng)用則等在基于離線生成性能預(yù)測的階段再估算。這樣既不犧牲吞吐又能盡早暴露深層的不一致。說到聯(lián)合調(diào)優(yōu)還有個細(xì)節(jié)一定要讓軟件團(tuán)隊在早期就參與驅(qū)動接口定義。很多芯片團(tuán)隊硬件做完才甩給軟件說“你們的編譯器適配吧”結(jié)果兩邊對寄存器接口的理解根本不在一頻道。我習(xí)慣做法是每個硬件模塊定義完之后立即讓軟件寫一個最小的煙霧測試smoke test在仿真環(huán)境里跑通基本算子有問題當(dāng)天反饋。這樣哪怕到流片時才暴露BUG也能迅速縮小到具體模塊而不是在全系統(tǒng)里盲人摸象。5. 從實測數(shù)據(jù)看軟硬件協(xié)同調(diào)優(yōu)的收益5.1 核心數(shù)據(jù)指標(biāo)怎么讀真正開始調(diào)優(yōu)了怎么看數(shù)據(jù)業(yè)內(nèi)有幾個關(guān)鍵指標(biāo)我每次都會讓團(tuán)隊把這幾項打印出來算子占用率、內(nèi)存帶寬利用率、片內(nèi)復(fù)用率、任務(wù)啟動開銷占比。算子占用率計算單元真正在算的時間占比。長期低于50%說明調(diào)度或切塊有問題。內(nèi)存帶寬利用率DDR或SRAM端口實際傳輸和理論峰值之比。高占用率但性能差可能是數(shù)據(jù)搬運(yùn)重復(fù)了而不是帶寬不夠。片內(nèi)復(fù)用率每次數(shù)據(jù)從內(nèi)存搬進(jìn)片上后被計算單元用了多少次。大于10基本健康小于3基本在發(fā)呆。任務(wù)啟動開銷占比驅(qū)動和運(yùn)行時消耗周期在所有計算周期中的比例。超過20%說明你得好好做算子融合了。以上這些指標(biāo)不要一次性全看那樣容易信息過載。我是習(xí)慣把一張熱力圖打出來橫軸是時間縱軸是不同硬件單元的活動度。哪一個單元一直在跑、哪一個單元閑得像條咸魚一眼就看出來。比如計算單元燒滿但數(shù)據(jù)傳輸線沒幾個脈沖那就說明數(shù)據(jù)沒喂飽反過來搬運(yùn)一直堵計算一停一停的就說明調(diào)度沒鋪開或者切塊切小了。5.2 一個真實的調(diào)優(yōu)實例我挑一個自己印象最深刻的小案例講吧。某圖像處理Demo輸入是一張1080p圖片跑一個簡單的超分模型總共就二十來個算子。上線后延遲慘不忍睹預(yù)估要50毫秒實際跑了200毫秒。最開始我們懷疑是卷積陣列太弱于是翻性能計數(shù)器結(jié)果發(fā)現(xiàn)卷積陣列占用率不到30%說明問題根本不在計算本身。進(jìn)一步看內(nèi)存事務(wù)的時間戳發(fā)現(xiàn)有很多頻繁的“小搬”操作每個算子都要讀一遍完整特征圖片內(nèi)復(fù)用率只有2。說白了數(shù)據(jù)根本沒有在片上停留下來每次計算都要從DDR把整個圖重新抬一遍。當(dāng)時我們的操作很簡單先把所有算子的數(shù)據(jù)布局統(tǒng)一成硬件友好的格式避免layout轉(zhuǎn)換產(chǎn)生的零散搬運(yùn)再把相鄰的卷積與加融合成一個算子讓中間特征圖停留在SRAM而不是一算完就落回DDR。幾天改下來延遲降到70毫秒后來又做了一次循環(huán)分塊細(xì)分把部分重復(fù)讀取進(jìn)一步削掉干到了45毫秒。整個過程中芯片的算力一點沒變變的全是軟硬件怎么協(xié)作、數(shù)據(jù)怎么流動。從那次之后我就很篤定一件事大多數(shù)AI芯片的性能事故不是算力不夠而是軟件沒有把數(shù)據(jù)喂好。硬件設(shè)計得再好數(shù)據(jù)流不轉(zhuǎn)就是空轉(zhuǎn)。6. 常見問題與排查技巧實錄6.1 排查思路先定位瓶頸在哪一級做軟硬件協(xié)同最怕的是拿到性能問題就到處懷疑。我習(xí)慣按一套固定排查流程走能省很多精神內(nèi)耗。先看計算單元占用率。如果占用率中等或偏上說明計算不是瓶頸趕緊去看存儲搬移和等待時間。如果占用率特別低先懷疑算子切塊是不是切得太大導(dǎo)致數(shù)據(jù)排不上再懷疑指令調(diào)度是不是串行依賴沒被打斷然后再懷疑數(shù)據(jù)布局是不是格式轉(zhuǎn)換把時間都燒走了。這套流程不算高深但真能攔住90%的初調(diào)問題。還有兩類坑很隱蔽。一類是同步陷阱多核芯片上如果各核心之間要頻繁同步任務(wù)等待時間會急劇膨脹。我見過一個模型算子都很輕但同步點特別密集導(dǎo)致多核利用率只有40%。后來把多個同步等待合并成一次握手一下子就提上來了。另一類是緩沖區(qū)分配錯誤有的編譯器會給每個算子開一塊固定大小的SRAM結(jié)果明明很多算子的生命周期沒重疊資源也沒共享白白浪費空間間接逼迫數(shù)據(jù)外溢到DDR。6.2 常見問題速查現(xiàn)象常見原因處理套路計算單元空閑內(nèi)存正常切塊尺寸不匹配片上數(shù)據(jù)復(fù)用低調(diào)整tile尺寸增大片內(nèi)駐留時間內(nèi)存帶寬打滿性能仍差數(shù)據(jù)被反復(fù)搬運(yùn)多次優(yōu)先做算子融合再統(tǒng)一layout單核快多核不提升同步點過密或任務(wù)分配不均合并同步、按數(shù)據(jù)并行重排任務(wù)小算子延遲極高驅(qū)動任務(wù)啟動開銷占比高算子融合或批量下發(fā)編譯器生成了冗余指令后端的模式匹配沒覆蓋硬件特性補(bǔ)上針對特定張量指令的匹配規(guī)則性能模型和實測偏離大模型里面存儲延遲估算不準(zhǔn)校準(zhǔn)各級存儲的實際cycle數(shù)6.3 兩個容易被忽略的細(xì)節(jié)最后說兩個我踩過坑的細(xì)節(jié)。第一個是“仿真器太快了”。我在早期調(diào)試時總覺得指令集模擬器跑得飛快程序一下子就出結(jié)果結(jié)果到了虛擬原型上一跑性能完全兩樣。原因是模擬器默認(rèn)所有訪存都命中而真實系統(tǒng)里緩存miss、SRAM bank沖突到處都是。所以建模擬器的時候一定要把存儲的bank沖突模擬進(jìn)來否則你辛辛苦苦調(diào)的指令序一到真硬件就垮。第二個細(xì)節(jié)是“性能數(shù)據(jù)要留到能溯源的程度”。很多芯片調(diào)試時計數(shù)器能告訴你總周期數(shù)但你看不出到底耗在哪個算子、哪條內(nèi)存邊上。所以我后來推行了一套規(guī)范每個算子執(zhí)行前記錄起始的流水線階段值執(zhí)行后讀一次算差值寫進(jìn)運(yùn)行日志。這樣做看似多余但真正排查問題時你才能在幾萬行日志里迅速定位到那一個操作出了毛病而不是拿著一個總時間反復(fù)猜。按照我個人的實踐體會AI芯片的軟硬件設(shè)計從來不是兩個聰明團(tuán)隊各寫各的再拼起來它是一場持續(xù)性的接口博弈。硬件留下多少余地軟件就能做出多少手腳軟件理解多少硬件硬件就能跑出多少性能。如果你正在做這顆芯片的工具鏈我建議從今天開始把每一個算子的性能計數(shù)器都打開把每一次數(shù)據(jù)搬移都畫出來。你很快會發(fā)現(xiàn)瓶頸永遠(yuǎn)比你想象得更靠前也永遠(yuǎn)更容易被一套嚴(yán)謹(jǐn)?shù)膮f(xié)同方法吃干榨凈。