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

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

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

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

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

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

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

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