AI編譯器的概覽、挑戰(zhàn)和實(shí)踐培訓(xùn)

AI編譯器的概覽、挑戰(zhàn)和實(shí)踐培訓(xùn)時間：2024年6月24日培訓(xùn)講師：XXX本次主要介紹AI編譯器的概覽AI編譯器的挑戰(zhàn)未來的展望Traditionalcompilervsaicompiler2024/6/22四大特征2024/6/22Python為主的動態(tài)解釋器語言前端多層IR設(shè)計(jì)：圖編譯器算子編譯器Codegen面向神經(jīng)網(wǎng)絡(luò)的特定優(yōu)化DSA芯片架構(gòu)的支持Python為主的動態(tài)解釋器語言前端

2024/6/22與傳統(tǒng)編譯器不同，AI編譯器通常不需要Lexer/Parser，而是基于前端語言（主要是Python）的AST將模型解析并構(gòu)造為計(jì)算圖IR，側(cè)重于保留shape、layout等Tensor計(jì)算特征信息，當(dāng)然部分編譯器還能保留控制流的信息。這里的難點(diǎn)在于，Python是一種靈活度極高的解釋執(zhí)行的語言，AI編譯器需要把它轉(zhuǎn)到靜態(tài)的IR上。多層IR設(shè)計(jì)2024/6/22為什么需要多層IR設(shè)計(jì)？因?yàn)樗梢赃_(dá)到易用性與高性能（一眼頂真，鑒定為鞭辟入里）為了讓開發(fā)者使用方便，框架前端會盡量對Tensor計(jì)算進(jìn)行抽象封裝，開發(fā)者只要關(guān)注邏輯意義上的模型和算子；而在后端算子性能優(yōu)化時，又可以打破算子的邊界，從更細(xì)粒度的循環(huán)調(diào)度等維度，結(jié)合不同的硬件特點(diǎn)完成優(yōu)化。因此，多層IR設(shè)計(jì)無疑是較好的選擇。GraphIR如MindSpore的MindCompiler（MindIR）、TF的XLA（HLO），TVM的Relay等，重點(diǎn)關(guān)注非循環(huán)相關(guān)的優(yōu)化。除了傳統(tǒng)編譯器中常見的常量折疊、代數(shù)化簡、公共子表達(dá)式等優(yōu)化外，還會完成Layout轉(zhuǎn)換，算子融合等優(yōu)化，通過分析和優(yōu)化現(xiàn)有網(wǎng)絡(luò)計(jì)算圖邏輯，對原有計(jì)算邏輯進(jìn)行拆分、重組、融合等操作，以減少算子執(zhí)行間隙的開銷并且提升設(shè)備計(jì)算資源利用率，從而實(shí)現(xiàn)網(wǎng)絡(luò)整體執(zhí)行時間的優(yōu)化。2024/6/22算子IR如MindSporeAKG、CANNTBE、TVM(HalideIR)等。針對算子層IR主要有循環(huán)變換、循環(huán)切分等調(diào)度相關(guān)的優(yōu)化，與硬件intrinsic映射、內(nèi)存分配等后端pass優(yōu)化。其中，當(dāng)前的自動調(diào)度優(yōu)化主要包含了基于搜索的自動調(diào)度優(yōu)化和基于polyhedral編譯技術(shù)的自動調(diào)度優(yōu)化（一眼頂真鑒定為純純的名詞黨）2024/6/22Codegen當(dāng)前基本上在LLVM上完成。還要提到的是MLIR，它實(shí)際上一種編譯的基礎(chǔ)設(shè)施。MLIR有Operation和Dialect的概念，Dialect，Operation，Attribute，Type等都可以通過td文件比較方便地定義出來。而LLVM定義新的intrinsic比較麻煩，定義新的IR就更麻煩了。LLVMIR主要表示硬件指令操作，而MLIR能表示更多東西，比如表示神經(jīng)網(wǎng)絡(luò)的圖結(jié)構(gòu)。因?yàn)橛蠨ialect，MLIR是組件化，去中心的，不像LLVM的ir是一種大而全的。MLIR執(zhí)行過程和LLVM一樣，IR會過由Pass組成的Pipeline，不斷地變換生成最終的IR。不同的是MLIR的IR可以是不同dialect的，構(gòu)成了Multi-Level的效果。2024/6/22面向神經(jīng)網(wǎng)絡(luò)的特定優(yōu)化2024/6/22自動微分：基于計(jì)算圖的自動微分、基于Tape和運(yùn)算符重載的自動微分方案、基于source2source的自動微分都是現(xiàn)在主流的方案。自動并行：隨著深度學(xué)習(xí)的模型規(guī)模越來越大，模型的并行優(yōu)化也成為編譯優(yōu)化的一部分，包括：數(shù)據(jù)并行、算子級模型并行、Pipeline模型并行、優(yōu)化器模型并行和重計(jì)算等（確實(shí)沒實(shí)操過任何一種方案，想有機(jī)會實(shí)操）

DSA芯片架構(gòu)的支持SIMT、SIMD、Dataflow：AI的訓(xùn)練和推理對性能和時延都非常敏感，所以大量使用加速器進(jìn)行計(jì)算，所以AI編譯器其實(shí)是以加速器為中心的編譯器，這個也是區(qū)別于通用編譯器的一個特征2024/6/22AI編譯器的發(fā)展歷程樸素AI編譯器框架和庫，動態(tài)圖和靜態(tài)圖圖和算子抽象：靜態(tài)圖采用了部分編譯器的技術(shù)（硬件無關(guān)優(yōu)化，e.g.表達(dá)式化簡、常量折疊）和硬件相關(guān)優(yōu)化（簡單算子融合）專用的AI編譯器Python原生表達(dá)的靜態(tài)化AI專用編譯器架構(gòu)，打開圖和算子的邊界進(jìn)行融合優(yōu)化通用化的AI編譯器圖算統(tǒng)一，融合優(yōu)化，算子自動生成更泛華的優(yōu)化能力：動態(tài)shape、稀疏、復(fù)數(shù)等自動并行優(yōu)化2024/6/22樸素的AI編譯器AI框架的早期主要是兩個抽象，一個是基于張量的計(jì)算圖，分為圖和算子；另外一個動態(tài)圖和靜態(tài)圖。動態(tài)圖基本上和AI編譯器沒有太多關(guān)聯(lián)，靜態(tài)圖采用了部分編譯器的思想，比如圖優(yōu)化的時候會使用一些硬件無關(guān)的優(yōu)化（表達(dá)式化簡/常量折疊等）、硬件相關(guān)的優(yōu)化（手工的算子融合）。問題：表達(dá)上：靜態(tài)圖的表達(dá)式非Python原生，開發(fā)者主要通過框架提供的API進(jìn)行顯式構(gòu)圖，易用性上不好；性能上：開發(fā)者定義的算子粒度未必最能發(fā)揮硬件的性能；硬件廠商的提供的缺省的算子庫也未必是最優(yōu)的，在模型和shape確定的情況下，可能還有更優(yōu)的算子實(shí)現(xiàn)；DSA芯片出現(xiàn)加劇了性能上的挑戰(zhàn)2024/6/22專用的AI編譯器主要的兩個特征：表達(dá)上，動態(tài)圖和靜態(tài)圖趨于一致，意味著AI編譯的入口更加與Python原生的表達(dá)接近；出現(xiàn)相對獨(dú)立的AI編譯器，聚焦在打開圖和算子邊界進(jìn)行融合優(yōu)化，發(fā)揮芯片的算力。問題：表達(dá)上，圖層和算子層的表達(dá)還是分開的，算法工程師主要關(guān)注圖層的表達(dá)，算子的表達(dá)和實(shí)現(xiàn)主要是框架開發(fā)者和芯片開發(fā)者來提供。功能泛化的問題：動靜轉(zhuǎn)換的成功率、動態(tài)shape、稀疏、分布式并行優(yōu)化等更多的需求無法滿足效率和性能的平衡：算子實(shí)現(xiàn)上在schedule、tiling、codegen上缺乏自動化手段，門檻高。2024/6/22通用的AI編譯器主要的特征：圖算統(tǒng)一表達(dá)；更泛化的優(yōu)化能力：動靜統(tǒng)一、動態(tài)shape、稀疏、復(fù)數(shù)、自動并行等；圖算融合優(yōu)化、算子自動生成?？偟膩碚f，個人感覺當(dāng)前的階段還是處于2.0階段，大家想盡快構(gòu)建通用AI編譯器的能力，但是還有許多關(guān)鍵的問題還沒有完全解決。（感覺這可能是以后一個收米的風(fēng)口，好日子還在后面呢.jpg）2024/6/22AI編譯器發(fā)展的驅(qū)動力和挑戰(zhàn)

Python的靜態(tài)化AI編譯器如何使能多樣性算力，特別是如何充分發(fā)揮DSA芯片的算力面向神經(jīng)網(wǎng)絡(luò)(NN)領(lǐng)域的特定優(yōu)化—自動并行+自動微分2024/6/22Python的靜態(tài)化Python靜態(tài)化是指通過JIT等技術(shù)，讓Python程序進(jìn)行靜態(tài)的編譯優(yōu)化，提升性能、方便部署，Python靜態(tài)化是AI編譯器開始工作的一個起點(diǎn)。業(yè)界Python靜態(tài)化的兩種方式：通用PythonJIT虛擬機(jī)：主要是期望在Python解釋執(zhí)行的基礎(chǔ)上增加JIT編譯加速的能力，典型的如PyPy；修飾符方式的PythonJIT方案：典型的如Numba，PythonJIT虛擬機(jī)的一種妥協(xié)實(shí)現(xiàn)方式，通過修飾符，進(jìn)行部分Python語句加速AI框架靜態(tài)化的方案普遍采用修飾符這套方案，這套方案細(xì)分下來也有兩種不同的方法：TracingBasedASTTransform2024/6/22Chanllenges類型推導(dǎo)：從Python動態(tài)類型到編譯器IR的靜態(tài)類型靈活的語法和數(shù)據(jù)類型轉(zhuǎn)換：slice、dict等控制流的處理JIT的編譯性能.....2024/6/22AI編譯器如何使能多樣性算力，特別是如何充分發(fā)揮DSA芯片的算力

前面提到AI的訓(xùn)練和推理都是對性能非常敏感的，所以在AI的場景中大量用到加速器，包括CPU的SIMD單元、GPU的SIMT架構(gòu)、NPU這樣的專用架構(gòu)等；特別是近期類Dataflow+SIMD這樣的DSA芯片占比逐步提升的情況下，AI編譯器逐步成為發(fā)揮這些多樣性算力的關(guān)鍵。統(tǒng)計(jì)數(shù)據(jù)表明，最近出現(xiàn)的AI芯片中，Dataflow架構(gòu)占比大于50%，這類架構(gòu)的特點(diǎn)：數(shù)據(jù)流圖的執(zhí)行調(diào)度更加能發(fā)揮芯片的性能，即芯片進(jìn)行整圖或者子圖的調(diào)度，而不是像GPU那樣，主流是kernelbykernel的調(diào)度較強(qiáng)大的Cube處理單元（類SIMD），較為復(fù)雜的內(nèi)存管理機(jī)制2024/6/22CHALLENGES性能優(yōu)化更加依賴圖算融合優(yōu)化，圖層和算子層獨(dú)立優(yōu)化無法充分發(fā)揮芯片性能，需要圖算融合優(yōu)化；子圖切分、子圖內(nèi)垂直融合優(yōu)化和水平并行優(yōu)化；優(yōu)化的復(fù)雜度提升，標(biāo)量+向量+張量+加速指令、多級的存儲結(jié)構(gòu)，導(dǎo)致Schedule、Tilling、Vectorization/Tensorization復(fù)雜。2024/6/22面向NN領(lǐng)域的特定優(yōu)化:自動并行+自動微分

自動并行依然是大模型訓(xùn)練的一個難題：當(dāng)前大模型訓(xùn)練碰到碰到的內(nèi)存墻、性能墻依賴復(fù)雜的切分策略來解決，包括：Scaleout：多維混合并行能力，含：數(shù)據(jù)并行、算子級模型并行、流水線并行、優(yōu)化器并行等Scaleup：重計(jì)算、Host/Devcie并行等這種方式最大的挑戰(zhàn)就是效率墻：如果依賴手工去配置切分策略，對算法工程師來說，門檻高，效率低；當(dāng)前類似半自動并行的方式可以解決一部分效率的問題，但是真正要解放工程師還是依賴編譯+尋優(yōu)結(jié)合，自動化的找到并行策略面向未來AI+科學(xué)計(jì)算場景，自動微分的要求更高，是另外一個挑戰(zhàn)這里有兩個大的難題：控制流：傳統(tǒng)的自動微分都是通過控制流展開方式來解決問題，動態(tài)圖通過正向在Python側(cè)執(zhí)行進(jìn)行控制流展開，一旦循環(huán)次數(shù)多的話，性能劣化；靜態(tài)圖的控制流自動微分目前還沒有太完善的方案。高級微分的性能：前向微分+后向微分；Jacobianmatrix(雅克比)；高階微分：Hessianmatrix（海森矩陣）2024/6/22未來一些開放性的問題

2024/6/22AI的圖層優(yōu)化最終是一個什么問題?整體而言，AI圖層的優(yōu)化當(dāng)前有兩種思路：編譯器的優(yōu)化思路，基于PatternMatch，從搜索算法的本質(zhì)講是一種貪心的算法，是局部尋優(yōu)的思路全局規(guī)劃的思路，無論是動態(tài)規(guī)劃、還是類似ILP的線性規(guī)劃也好，都是采用全局尋優(yōu)的思路，這里的關(guān)鍵是需要一個較好的costmodel。未來圖層的優(yōu)化哪一種方式是主流，或者哪幾種方式組合最優(yōu)。2024/6/22圖算能否統(tǒng)一表達(dá)，統(tǒng)一編譯優(yōu)化