AICompiler编译器介绍及访存密集算子优化
阿里云機器學習PAI是阿里云上提供的人工智能平臺服務，提供一站式的機器學習解決方案。在過去的三年多時間里，阿里云機器學習PAI團隊在AI編譯優化技術方向投入了比較專注的資源精力，對這個領域也建立起了一定的理解。目前在內部集群上AICompiler已經實現了默認的全量打開，在阿里外部已經開始為來自多個不同行業的業務方提供服務，涉及行業領域包括在線教育，安全，娛樂，電商，直播等；涉及的業務類型包括文本分類/識別，機器翻譯，語音識別/生成，圖像類業務等等；支持覆蓋的硬件包括云端GPU/CPU以及端上CPU等。
我們會在接下來的系列文章里介紹一下這方面的工作以及一些思考，一方面拋磚引玉，吸引更多同學來進行共同建設；另一方面也為阿里云上對訓練和推理作業的性能優化有一定需求的用戶同學提供一些技術方面的背景介紹。
AI編譯器簡介
在上圖中將近年的深度學習框架粗略分為三代。幾代框架之間有一個趨勢，在上層的用戶API層面，這些框架在變得越來越靈活，靈活性變強的同時也為底層性能問題提出了更大的挑戰。另外一個技術趨勢則是系統底層的深度學習的編譯器近一段時間也開始活躍起來，這些編譯器試圖去解決框架的靈活性和性能之間的矛盾。
傳統編譯器是以高層語言作為輸入，避免用戶直接去寫機器碼，而用相對靈活高效的語言來工作，而深度學習編譯器的作用相仿，其輸入是比較靈活的，具備較高抽象度的計算圖，輸出包括CPU或者GPU等硬件平臺上的底層機器碼及執行引擎。
AI編譯器的目標是針對AI計算任務，以通用編譯器的方式完成性能優化。讓用戶可以專注于上層模型開發，降低用戶手工優化性能的人力開發成本，進一步壓榨硬件性能空間。
涉及到性能優化，我們有必要先對一個AI作業執行過程中的性能開銷的分布有一個感性的認識，所謂what you can't measure, you can't optimize it. 在《Characterizing Deep Learning Training Workloads on Alibaba-PAI》這篇paper里，我們針對阿里云機器學習PAI平臺訓練workload的性能開銷占比有過一個比較細致的分析。考慮到目前我們在AI編譯優化里還主要關注單計算設備的計算效率問題，所以我們可以宏觀上將單設備上的性能開銷拆解為計算密集算子（比如GEMM和Convolution）和訪存密集算子（比如Elementwise Add，BN等操作）兩部分，關于計算圖過于靈活帶來的框架開銷，我們也統一歸類到訪存密集算子開銷占比中。
針對不同的性能熱點，所需要的優化手段也存在區別。本文后續篇幅將首先介紹我們在訪存密集型部分的工作，計算密集型算子的部分會在后續連載中進行介紹。
Fusion Stitching——大顆粒度訪存密集型子圖CodeGen
阿里云機器學習PAI團隊在訪存密集算子上以XLA為基點開展了大量的工作，工作內容涵蓋前端、中端、后端以及執行引擎，我們將其稱之為Tensor Accelerator and Optimizer(TAO) compiler。
TAO Compiler先后經歷了V1 (FusionStitching: Deep Fusion and Code Generation for TensorFlow Computations on GPUs)和V2(FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads)兩代演進。
如前文所說，XLA在GPU backend上的主要收益來源是對訪存密集型算子的自動Op Fusion CodeGen。催生我們做這些嘗試的原因是，我們在實際業務中發現，社區XLA在最核心的CodeGen環節還有很大的問題和改進空間。
例如下圖為一個LayerNorm模塊的前向計算子圖，手工優化的話，它可以很容易被寫成一個CUDA kernel，本應該是很適合編譯器通過自動op fusion來獲取性能收益。但我們發現XLA實際并沒有做的非常好。經過細致的分析后我們認為這里的本質問題是社區的CodeGen模版過于簡單，因為模版非常簡單限制了中端Fusion pass的靈活度，被迫做了非常保守的op fusion策略，圖中每一條紅線都有一個細節的原因導致社區XLA無法把這些算子fuse到一起，也就無法拿到節省相應訪存開銷所能得到的優化收益。
分析之后我們認為造成這些問題的根本原因為：
? 單一Parallel Loop的比較簡單的CodeGen模版；
? 由于CodeGen模版比較簡單，為保證性能而被迫保守的Op Fusion策略。
這種簡單的CodeGen策略在fuse線性的Elementwise算子時足夠得到理想的結果，但當計算圖的連接關系變復雜，同時計算圖中出現Reduce／Broadcast等復雜計算節點時，在實際業務中效果并不好，特別是對于包含了后向處理部分的訓練過程計算圖更是如此。
當然，我們可以直接采用復雜的CodeGen模版，通過犧牲通用性的方式換取更好的性能收益。但在Op節點類型以及計算圖的拓撲結構千變萬化的情況下，這種方式并不能夠在編譯器中通過編譯的方式有效的節省人力成本。于是就引申出了TAO compiler的設計理念。
TAO Compiler V1
TAO Compiler V1的核心原理可以概括為：
? 借助于GPU硬件中低訪存開銷的shared memory，縫合不同schedule的計算子圖，實現多個parallel loop復合
? 支持不同shape的多個tensor輸出
? 保證CodeGen性能的前提下，實現更加激進的Op fusion策略
在GPU的體系結構中，SM內的shared memory的訪存開銷遠遠低于global memory的訪存開銷，可以幫助我們來把多個本來獨立的kernel粘合在一起。這樣的方案下，被粘和的每一個子圖仍然可以基于自己的簡單的CodeGen模版來運作，他們不必擔心因為各自的pattern被fuse在一起而需要更復雜的CodeGen模版，也一定程度上不必擔心因為fuse的顆粒度太大而產生過多的冗余的計算。V1的核心工作涉及中端和后端兩個部分，解決了多種不同細節原因導致的細粒度算子無法fuse到一起的問題，主要是基于這個核心思路。
下面的兩張圖分別為社區XLA Codegen的情況和TAO Compiler V1的情況。TAO Compiler的CodeGen將社區做法中難以fuse在一起的計算pattern，通過低訪存開銷的shared memory作為橋接，將多個kernel縫合在一起，同時又保留被縫合的每一個部分都依然采用簡單的CodeGen模版。這種方案在通用性和性能之間取得一個更好的平衡點。
TAO Compiler V1通過這樣的方式，在一定程度上保證CodeGen性能的基礎上，大幅提高了op fusion的顆粒度，以下面表格中的兩個例子為例：
TAO Compiler V2
在做V1的推廣落地的同時，我們抽取了一些業務進行了比較精細的分析，觀察優化距離性能的極限還有多大的距離，結論是雖然憑借在op fusion顆粒度上的優勢，在很多case上可以做到明顯優于社區XLA，但離硬件性能的峰值還是有比較明顯的gap的。所以我們抽取了一些有代表性的kernel，分析了這些gap的具體原因。我們發現每個case各自都有不同的原因導致沒有能夠接近訪存的峰值性能。這些原因比較細節，在此不再展開，但宏觀來說問題來自于目前的CodeGen基于Rule-Based的策略，而Rule-Based的策略本身存在其固有的局限性。
具體到compiler的各個環節：
? 在中端Op Fusion層面上，基于簡單規則下局部合理的Op Fusion決策可能并非全局最優的決策；
? 在后端CodeGen層面上，基于簡單規則下，同一Kernel內的不同節點的模版選擇和CodeGen行為決策也可能并非整體最優的決策。
在TAO Compiler V2與V1以及社區XLA相比最大的區別可以概括為：
? 更多的考慮全局最優，而非基于貪婪的局部最優；
? Rule-Based到Cost-Based的演進。
這兩點區別同時反映在中端的Op Fusion策略以及后端的CodeGen策略上。
我們引入了多個不同等級的Cost Model，其主要區別如下：
中端在Graph范圍內做Op Fusion決策，由于探索范圍較大，選擇接入開銷較低的SOL cost model；后端在Kernel范圍內做tuning，可以接入開銷相對較高的Proj Cost Model。
中端Op Fusion決策部分
目前XLA和TAO V1中解決HLO instruction之間的Auto Fusion問題時采用的是完全基于規則的策略。這些規則都是基于人工經驗確定的，往往難以覆蓋到所有的情況，導致可能會出現各種corner case以及為了解決這些corner case而打上的一系列補丁。從長期來看這種方法并不能充分的挖缺潛在的性能優化空間，也會帶來越來越沉重的維護成本。為了使用更加系統性的方法解決Auto Fusion的問題，我們提出了Cost Based Fusion Strategy。
新的策略主要是對V1中以下兩個大的方面進行改進：
? 兼容性度量。 兼容性度量是一種用來評估一個給定的Fusion動作是否有收益的方法。目前V1中使用的是一條簡單的規則來進行兼容性度量，即只要融合之后的kernel的thread block level的并行度沒有下降到1則認為融合具有收益。該策略在很多實際業務上發揮了很好的作用，但也遇到了不少Corner case。比如該規則對于大shape的pattern表現為過于激進，融合后可能因并行性大幅下降而導致性能變差，而對于小shape的pattern則表現的過于保守，丟掉了一部分融合的機會。在V2中我們基于cost model來更系統化的考慮訪存量，kernel launch次數，并行度改變對于最終性能的影響等，可以更準確的評估一個融合動作的效果以指導Fusion Plan的探索。
? Fusion的探索空間。 一個計算圖對應的Fusion Plan由多個Fusion Pattern組成。一個Fusion Pattern中如果只包含具有生產消費關系指令則是縱向fusion，反之如果包含有不含生產消費關系的指令否則為橫向fusion。一個計算圖可能存在著多種不同的Fusion Plan。目前V1中是基于一個啟發式的貪心策略確定性的生成一個Fusion Plan，而再不進行其他可能的探索。這種策略好處是具有很高的效率，但缺點是會漏掉很多可能性。另一方面目前社區XLA及TAO Compiler V1中只會考慮縱向fusion而未考慮橫向fusion。這里所說的橫向fusion即不存在直接生產消費關系的節點之間的fusion，如下圖所示。除其中近距離橫向fusion可同樣帶來仿存節省的收益外，在更多的場景下可通過節省kernel launch開銷帶來收益，同時當使用空分并行的方式對橫向fusion之后的kernel做CodeGen時還可能提升GPU的資源利用率。下圖中展示了一些橫向fusion的例子。V2的中端部分，對包含以上兩種fusion的解空間進行了更多探索，以期找到更好的全局最優決策。
具體來說fusion空間探索是一個迭代式的過程，每一次迭代由以下幾個步驟構成：
? 搜索高價值fusion pattern。 這里的高價值指的是基于cost model評估會帶來更大收益的pattern。直接暴力枚舉進行搜索是一個指數級別的算法，我們使用的是一種啟發式的算法，在稀疏圖場景下（一般指令間的連接圖滿足這個前提）是一個接近線性的算法。具體而言我們首先構造計算圖指令間的一個拓撲排序，依次遍歷每條指令并按照遞推的方式生成以該指令結尾的top k候選fusion pattern。遞推的過程中我們會進行局部的充分枚舉（擴張），對整個空間進行更多的探索，同時也會及時進行剪枝（收縮），確保可以不會組合爆炸。
? 構造fusion plan。 即尋找Fusion Pattern的一個有效組合以構成一個完整的Fusion Plan。這里的有效組合指的是各個Fusion Pattern之間不相交且相應的融合動作不會引入環。 這里我們嘗試了多種不同的策略，比如貪心策略，BFS+剪枝策略。目前使用的是每次選擇兼容的且收益最大的fusion pattern的貪心策略。
? apply fusion plan。 也即根據Fusion Plan執行相應的融合動作。
以上的過程可以對整個優化空間進行更大程度的探索，同時又有效的控制了搜索的復雜性，在一組評測應用集合（這也是進行編譯器開發經常性需要的工具支撐了）上的實驗表明，該策略工作良好，符合預期。
后端CodeGen部分
在介紹V2后端CodeGen的具體工作之前，我們先對XLA和TVM做一個簡單的橫向比較：
? XLA是一種比較務實的方案，用基于規則的簡單方式去做收益空間最大的訪存密集型pattern的Automatic CodeGen，它的優勢是自動化和比較小的編譯開銷，缺點是如果想進一步提升性能的話會非常局限；
? TVM的目標是去勝任更復雜的CodeGen，基于規則的策略無法勝任情況下，它在Halide工作的基礎上，提出了一個比較先進的Schedule抽象。這個Schedule的空間非常大，如何才能找到最好的Schedule，一方面需要基于用戶的經驗寫出比較適合硬件體系結構的Schedule，一方面需要巨大的編譯開銷為代價在Schedule之間進行tuning。
這里我們需要解決的問題是：需要找到一種方式，能夠自動/對用戶透明的，在有限的編譯開銷的基礎上，提升性能并進一步接近硬件的性能峰值極限。
V2的工作在XLA的基礎上，一定程度上試圖借鑒TVM的Schedule抽象的核心想法，但又有本質的不同。根據我們過往工作的經驗，對于訪存密集型pattern來說，為了達到或接近硬件的峰值性能極限，所需要的Schedule的枚舉空間通常并不需要非常大。因此我們試圖提煉一套比TVM的schedule抽象要簡化很多的CodeGen Plan抽象，希望這套抽象能夠代表手工寫相應Pattern的Kernel時所需要考慮的全部問題空間。同時需要保證CodeGenPlan的枚舉空間是足夠小的，能夠滿足用戶對編譯開銷的需求。此外，我們使用Proj Cost Model而非TVM中目前基于實際profiling結果來tuning的方式，其原因也是希望能在編譯開銷和性能之間取得更好的平衡點。
在TAO Compiler V1的工作的基礎上，V2在CodeGen部分的主要區別是：
? CodeGen Plan抽象，盡量對全部的CodeGen可能性做抽象和枚舉；
? Index計算與主體計算分離，引入index_cache與value_cache，最小化冗余計算；
? 與CodeGen Plan抽象相對應的Implementation抽象，方便對不同節點增加新的算子實現模版。
篇幅原因我們不再詳細介紹CodeGen Plan抽象的具體細節。下面僅試圖以一個pattern為例，來說明TAO Compiler V2與V1及社區XLA之間，在CodeGen結果上的主要區別。
下圖左圖為社區XLA的CodeGen結果，由于簡單模版的限制會生成3個簡單Schedule的kernel；右圖是TAO Compiler V1的結果，借助Shared Memory的粘合，可以生成一個kernel。
TAO Compiler V1距離硬件性能峰值之間可能存在差距的主要問題來自于：
? A/B/C三個子部分的CodeGen Schedule是彼此獨立的，在此例中，A/B/C三個部分的計算在同一個GPU thread內可能依賴于Elementwise 4節點處的不同Element，產生冗余計算。
? A/B/C三個子部分的CodeGen過程相對獨立，Elementwise 3等節點處的值和Index的計算結果無法被有效Cache，進一步增加了冗余計算指令。
? A/B/C三個子部分可能各自有多種CodeGen模版，但可能只有一部分子模版的組合可以在全局上最好的利用GPU的并發度（即GPU Occupancy）。
當冗余計算指令的數量（或者Occupancy帶來的問題）達到一定臨界值后，本應該是仿存Bound的計算pattern，性能熱點便會遷移為計算或Occupancy，導致無法達到硬件仿存性能峰值的性能。
TAO Compiler V2經過CodeGen Plan的枚舉和基于Cost Model的tuning之后，會找到一個認為相對最優的CodeGen Plan。例如，在上圖這個pattern中，選擇Reduce節點做為整個Kernel的Dominant節點可能是更好的方案，其它節點處，在同一個thread內需要計算的Index會由Dominant節點的Index以輻射傳播的方式推導計算得到，Index的傳播過程和數值的計算過程在整個kernel范圍內進行Cache，通過以上機制最大化的避免冗余的Index計算和數值計算。此外，CostModel還會對枚舉的CodeGen Plan中的Launch Dimension等因素對Occupancy的影響進行建模，得到對于整體性能最優的CodeGen行為。
TAO Cost Model
中端不同的fusion策略以及不同后端CodeGen生成方式會影響最終模型實現的效率。所以很自然的想到利用cost model幫助中端、后端CodeGen探索最優的實現。這里主要的挑戰是：
? cost model需要高效的對成百上千種不同實現做實時的評估，需要減少其本身的overhead；
? cost model需要準確的區分不同實現的性能差異；
? 不同型號的gpu擁有不同的硬件架構（memory hierarchy，latency以及throughput指標），cost model需要對不同型號的gpu建模。
所以我們會將workload輸入映射成不同的資源需求，同時對device內部流水線stage的resource、latency以及throughput做了抽象和表示。最終通過建立latency模型以及throughput模型估計當前實現的計算耗時。
Cost Model分為3個level： SOL Cost Model, Projected Cost Model以及Profiling Cost Model。
? SOL Cost Model 只考慮硬件spec，不考慮具體實現、架構特點（例如假設100% cache hit ）和編譯器效率。對于CodeGen，主要用memory cycles以及device throughput（例如：處理 1000T MACs, V100 125 Tflops vs 4x TPU 180 Tflops vs new ASIC）來區分。其特點是建模速度快，誤差較大， 所以可以用于粗粒度性能篩選以及不同硬件性能快速評估。因此,我們在中端OP fusion探索決策時引入SOL cost model評估不同Op fusion的cost并保留下cost 最小的Op fusion Plan。
? Projected Cost Model需要考慮具體kernel實現、架構特點（如hit rate）和編譯器特點。對于CodeGen，主要用kernel latency來區分。特點是建模速度慢，但誤差較小，所以可以用于細粒度性能篩選。因此,我們在做后端CodeGen時，引入Proj cost model區分同一條fused instrcution在不同CodeGen Plan下因為指令數不同以及寄存器復用情況差異造成的性能差別。
? Profiling Cost Model目的是作為cost model的ground-truth， 從profiler上得到實際運行性能結果。其特點是執行耗時長，所以可以用于構造offline cost model以及修正SOL level與Proj Level cost model。
Cost Model主要由三個模塊組成： Workload Model, Device Model, 以及 Cost Evaluation。簡單來說 。可以舉一個簡單的例子，假設任務的workload由10條硬件指令組成，芯片的處理能力是每個cycle處理2 條硬件指令，那么該workload需要5個cycle才能被處理完。
Workload model的準確性影響了整個cost model建模的準確性，因此80%以上的研發時間用于建模workload model。實踐中，我們會估計某條fused instruction在當前CodeGen Plan下的寄存器消耗、各個硬件單元的硬件指令需求、memory 消耗以及shared memory消耗。
Device Model建模了GPU各個硬件單元的throughput、bandwidth以及latency。可以從公開資料或者論文中得到.在device model中，我們主要建模了芯片內部各個unit的數量、SM內各個單元的throughput和latency、L2的throughput（用于建模atom指令）以及 Memory bandwidth。
Cost Evaluation根據workload model以及device model 估計最終生成kernel的Memory cycle、Occupancy（并行度）、Wave number、Kernel Latency、Performance Limiter以及SM內部各個unit的cycles。
Cost Model主要有兩種建模方式：Throughput Model以及Latency Model。
Throughput Model首先估計fused instruction所需各機器指令的數量（例如需要X條FFMA, Y條MOV, Z條IMAD等）。并對機器指令根據其對應的pipe做聚合（例如FFMA和IMAD都是發射到FMA Pipe的）。在獲得當前GP各個pipe的throughput后，用pipe workload（硬件指令）的數量除以這個pipe的throughput就是其所需要的cycle數。最后在所有pipe中找cycyle最大的pipe，這個pipe就是這個fused instruction在當前gpu上的performance limiter，并且這個pipe所需要的時間就是kernel的時間。
Latency Model首先估計fused instruction所需的register數量，并計算出它的occupancy，再計算出總共需要的wave數量（波數）。然后估計fused instruction所需各硬件指令的數量（同Throughput Model中的第一條）。根據device model獲得當前gpu各個硬件指令的latency。根據硬件指令數以及指令的latency計算 warp latency，用warp latency乘以wave數量得到最后的kernel latency。
對于TAO Compiler V1來說，由于大部分fused instruction已經達到或者接近GPU math或memory的理論上限，此類情況下，throughput model能很好地對fused instruction做建模。 對于TAO Compiler V2，由于fuse了更大的尺度，目前大部分fused instruction距離硬件極限還有點距離，所以用throughput model做區分是不合理的。需要用latency model。
性能數字
下圖為基于本文中所述工作，在一些典型業務模型上，TAO Compiler與社區XLA在Nvidia V100上的性能數字對比，其中社區XLA與TAO的性能收益數字，均為相同計算精度下的收益數字。
下圖為某業務方基于PAI的easyTransfer解決方案進行建模，在不同模型上的性能收益數字。AICompiler同樣對不同應用場景給出了一致性的優化方案和相對比較明顯的性能優化結果。提供的性能優化數字普遍超過了社區XLA的工作。
阿里云機器學習PAI平臺面向企業客戶及開發者，提供輕量化、高性價比的云原生機器學習平臺，涵蓋交互式建模、拖拽式可視化建模、分布式訓練到模型在線部署的全流程覆蓋，百余種落地場景，全面提升機器學習工程效率。
目前， PAI AICompiler已經集成在阿里云機器學習PAI的通用推理優化工具Blade敏捷版中，用戶可以參考開發文檔來快速體驗。
作者：朱凱，趙文益，楊軍
原文鏈接：https://developer.aliyun.com/article/782717?
版權聲明：本文內容由阿里云實名注冊用戶自發貢獻，版權歸原作者所有，阿里云開發者社區不擁有其著作權，亦不承擔相應法律責任。具體規則請查看《阿里云開發者社區用戶服務協議》和《阿里云開發者社區知識產權保護指引》。如果您發現本社區中有涉嫌抄襲的內容，填寫侵權投訴表單進行舉報，一經查實，本社區將立刻刪除涉嫌侵權內容。 與50位技術專家面對面20年技術見證，附贈技術全景圖總結
以上是生活随笔為你收集整理的AICompiler编译器介绍及访存密集算子优化的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 阿里云 OpenAPI 开发者门户全新上
- 下一篇： 新手必看 | RVB2601开发板快速上