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

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

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

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

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

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

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

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