基于推導的人工智能編譯器-v3-翟季東.pdf

《基于推導的人工智能編譯器-v3-翟季東.pdf》由會員分享，可在線閱讀，更多相關《基于推導的人工智能編譯器-v3-翟季東.pdf（37頁珍藏版）》請在三個皮匠報告上搜索。

1、清華大學 Tsinghua UniversityPACMAN基于表達式推導的人工智能編譯器基于表達式推導的人工智能編譯器翟季冬翟季冬清華大學清華大學2023.3.26第十屆開源操作系統年度技術會議第十屆開源操作系統年度技術會議-OS2ATCOS2ATC清華大學 Tsinghua University背景和意義人工智能芯片人工智能芯片成為深度學習應用的主要計算力成為深度學習應用的主要計算力摩爾定律發展逐漸放緩深度學習對計算巨大需求為了降低編程復雜度，為了降低編程復雜度，深度學習編程框架深度學習編程框架及及人工智能編譯器人工智能編譯器應運而生應運而生人工智能芯片人工智能芯片Caffe編程框架編程框

2、架?人工智能編譯器人工智能編譯器2清華大學 Tsinghua University研究背景和意義深度學習模型的編譯過程深度學習模型的編譯過程深度學習代碼深度學習代碼人工智能芯片人工智能芯片體系結構無關無關的圖層優化(圖融合、替換、刪冗)中間數據流圖中間數據流圖體系結構相關相關的算子優化數學算子庫，數學算子庫，張量編譯等張量編譯等為了提高人工智能應用的優化與部署效率，現有的張量編譯器將優化過程解耦為高層次的圖優化和低層次的算子優化3清華大學 Tsinghua University研究現狀-1：圖層優化傳統的優化框架以基于規則的圖優化為主，如 TensorFlow中的 Grappler，需要領域專

3、家對優化規則進行人工設計為了提高優化效率，以 TASO、PET 為代表的自動優化框架成為當前圖優化領域發展的主流趨勢優點可自動地對圖優化空間進行搜索利用形式化驗證保證正確性缺點優化結果只能包含算子庫中的算子需嘗試大量錯誤狀態，優化時間長4清華大學 Tsinghua University研究現狀-2：算子層優化以 cuDNN、cuBLAS、oneDNN 等為代表的人工智能算子庫在特定硬件平臺上提供極高的性能，但需要大量人力進行維護近年來以 TVM、Ansor 為代表的算子自動生成框架成為了發展趨勢優點基于原語的編程接口相對簡單具有一定的跨平臺能力缺點很難達到算子庫的性能缺少對圖層信息的感知行優先

4、列優先并行分塊向量化并行分塊典型的調度原語a=tvm.nd.array(numpy.random.rand(M,K).astype(dtype),ctx)b=tvm.nd.array(numpy.random.rand(K,N).astype(dtype),ctx)k=te.reduce_axis(0,K),k)A=te.placeholder(M,K),name=A)B=te.placeholder(K,N),name=B)C=pute(M,N),lambda x,y:te.sum(Ax,k*Bk,y,axis=k),name=C)s=te.create_schedule(C.op)xo,y

5、o,xi,yi=sC.tile(C.op.axis0,C.op.axis1,bn,bn)(k,)=sC.op.reduce_axisko,ki=sC.split(k,factor=4)sC.reorder(xo,yo,ko,ki,xi,yi)sC.vectorize(yi)spackedB.parallel(x)用戶手寫的調度代碼缺少一套可以融合圖層與算子層優化的方法5清華大學 Tsinghua University現有圖層優化編譯器-TASOTASO：自動的等價計算圖搜索和替換搜索由后端支持的算子所組成的計算圖利用形式化驗證形式化驗證證明計算圖等價性用等價的高效高效計算圖替換低效低效計算圖計

6、算圖計算圖搜索搜索輸入子程序利用所有可用的算子，遍歷所有算子個數不大于某個固定閾值的計算圖后端支持的算子發現并利用等價的高效計算圖6清華大學 Tsinghua University基于局部等價優化的人工智能應用編譯優化Pet:Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections通過構造算子庫中沒有的內存布局重排算子，引入局部等價優化，為張量程序的優化帶來新的機會7清華大學 Tsinghua University局部等價優化8.=ConvW1W2XConvAdd

7、ConvW1W2XYZAdd完全等價優化缺點：會錯過一些優化機會優點：保證功能上的完全等價性張量中有成千上萬個元張量中有成千上萬個元素，是否可以在優化時素，是否可以在優化時只考慮“大部分”元素只考慮“大部分”元素的等價性？的等價性？局部等價優化局部等價優化清華大學 Tsinghua University局部等價優化9.=局部等價優化.XDilated ConvYWXConvZWConvW1W2XConvAddConvW1W2XYZAdd完全等價優化缺點：會錯過一些優化機會優點：保證功能上的完全等價性優點：可以取得更好的性能 更高效的算子 更高效的張量內存布局 更多的系統優化機會缺點：可能造成精

8、度降低、功能不等價清華大學 Tsinghua University局部等價優化10.=局部等價優化.XDilated ConvYWXConvZWConvW1W2XConvAddConvW1W2XYZAdd完全等價優化缺點：會錯過一些優化機會優點：保證功能上的完全等價性優點：可以取得更好的性能 更高效的算子 更高效的張量內存布局 更多的系統優化機會缺點：可能造成精度降低、功能不等價利用局部等價變形所帶來的性能提升的同時利用局部等價變形所帶來的性能提升的同時維持功能的等價性維持功能的等價性清華大學 Tsinghua University突變生成器11突變生成器突變糾錯器程序優化器突變突變生成器生成

9、器輸入子程序利用所有可用的算子，遍歷所有算子個數不大于某個固定閾值的子程序可發現完全等價完全等價與局部等價局部等價的優化變形后端支持的所有算子內存布局重排算子清華大學 Tsinghua UniversityPET 整體架構12優化后的程序突變突變生成器生成器突變程序突變突變糾錯器糾錯器完成糾錯的突變程序程序優化器優化器輸入程序清華大學 Tsinghua University程序優化器13基于搜索的基于搜索的程序優化器程序優化器輸入程序優化后的程序突變生成器與突變生成器與突變糾錯器突變糾錯器多重線性張量程序已完成糾錯的突變利用非線性算子非線性算子將程序劃分為子程序對每個子程序搜索其候選突變通過后

10、處理優化后處理優化來保證優化后程序的高效性突變生成器突變糾錯器程序優化器清華大學 Tsinghua University挑戰：如何檢測和糾錯？14利用多重線性張量程序良好的代數性質，快速定位不等價結果并生成糾錯內核Program fProgram g突變生成器突變糾錯器程序優化器具體見 PET:Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections.OSDI.2021.清華大學 Tsinghua University優化案例 dilated conv下圖中的例子

11、通過將輸入張量進行變形把 dilated conv 的計算變為普通conv 的計算15清華大學 Tsinghua University實驗結果在 batch size 分別為 1 和 16 的情況下，在常用的 DNN 模型上可取得最高最高 2.5 2.5 倍倍的加速比16清華大學 Tsinghua University小結：基于內存布局重排的局部等價優化提出了基于局部等價變換的張量程序編譯優化技術結合張量程序應用特征，創新性地提出了張量重排與局部等價的優化策略對常見的模型和算子可以有最高2.5倍的性能提升但該方法依然存在一定局限性優化空間受限：僅能利用用戶指定的算子進行搜索和優化無效搜索較多：

12、遍歷搜索需嘗試大量錯誤狀態，增加優化開銷17清華大學 Tsinghua University基于表達式推導的圖算融合優化EinNet:Optimizing Tensor Programs with Derivation-Based Transformations以表達式替代算子表示，利用代數運算規則對表達式進行自動推導，以發掘更多的圖算融合優化潛力18清華大學 Tsinghua University基于張量表達式的形式表達張量表達式可表示遍歷順序、維度取值空間、內存布局等張量程序的重要特征人工智能應用中常用的算子均可表示，如卷積、矩陣乘等卷積算子的張量表達式示例引入了循環符號，顯式表示出元素遍

13、歷的順序及取值空間顯式標示出循環和求和變量的取值空間，以支持取值空間的映射變換以索引方式表示內存布局，以支持不同的內存布局變量在保證完備性的同時，準確表示了張量程序的特征19清華大學 Tsinghua University在表達式中融合圖層信息通過表達式嵌套的方法將圖層信息融合到表達式中AW1BConvW2ConvAdd 將不同算子以嵌套方式放入同一個表達式，不同子表達式間可以進行聯合優化20清華大學 Tsinghua University利用表達式替代基于計算圖的算子表示AKBConv缺點：優化受限于算子庫 只能進行圖層優化 表示方法簡潔 底層硬件無關優點：缺點：缺少現有優化技術的支持 優化

14、的搜索空間變大 保留圖層信息的同時包含算子計算邏輯 底層硬件無關的同時擺脫算子庫束縛 可支撐圖算融合優化優點：基于算子的計算圖表示基于表達式的表示21清華大學 Tsinghua University表達式推導及算子映射的基本流程將已有算子映射到算子庫將已有算子映射到算子庫將算子庫沒有的算子進將算子庫沒有的算子進行自動生成行自動生成推導過程利用基本推導過程利用基本代數運算規則，如代數運算規則，如交換律、結合律等交換律、結合律等輸入輸入計算圖計算圖張量張量表達式表達式輸出輸出計算圖計算圖22清華大學 Tsinghua University優點：擴展優化空間相對現有自動圖優化工作的優化空間局限性僅能

15、利用用戶預定義的算子用戶預定義的算子進行搜索和優化局限于POR變換(predefined operator representable)利用張量表達式將優化空間擴展至非預定義算子探索涵蓋任意張量表達式的變換融合圖層和算子層信息，發掘更多優化機會23清華大學 Tsinghua University推導規則匯總推導規則推導規則形式化描述形式化描述*描述描述算子間算子間規則規則表達式拆分將一個表達式拆分為多個無依賴的表達式表達式合并將多個無依賴的表達式合并為一個表達式表達式融合將多個有依賴的表達式合并為一個表達式算子內算子內規則規則求和拆分將一個求和作用域拆分為兩個變量替換將遍歷操作的迭代變量進行替

16、換遍歷合并將兩個作用域下的遍歷操作合并為一個邊界放松將迭代變量的取值范圍擴大邊界收緊將迭代變量的取值范圍縮小推導規則可方便地進行擴展，以支持更多類型的優化*由于篇幅原因，此處只保留了推導規則的形式化描述的公式，省略了其限制條件24清華大學 Tsinghua University自動推導過程利用 TVM 進行代碼生成 代碼生成部分以訪存代碼生成部分以訪存密集型算子為主，復密集型算子為主，復雜度較低，調度簡單，雜度較低，調度簡單，搜索時間較短搜索時間較短 本研究的形式化表示本研究的形式化表示方法與方法與TVM的表達式的表達式有兼容，可以方便地有兼容，可以方便地通過代碼映射生成通過代碼映射生成TVM

17、代碼代碼25清華大學 Tsinghua University新變換的機會單個卷積算子的部分推導方案Conv3x3計算可變換為Matmul、Conv1x1、Conv2x2、Group conv等計算轉換為表達式(c)(b)().VSVS+TM+T0WMatmulOffsetConcat(i).()(g)(h)求和拆分變量替換SSVS推導規則遍歷合并邊界放松邊界收緊TMBRBTeOpOp預定義算子張量表達式算子沿循環變量求和公式推導表達式實例化()Conv2x2T0WOffsetReducer2s2Pad+Split(n)BRConv2x2T0WOffsetAddSplitConv2x1 Conv

18、1x2 Conv1x1(p)(22)(11)(11)(m)(22)(11)(11)(o)T0WGroup Conv3x3Duplicate(k)BRVS+SS+(c)(l)VS+TM+()(j)fMatmulT0WOffsetReducers(e)可行變換方案T0W(a)Conv3x3輸入程序Conv1x1T0WOffsetReducekTranspose(f)SSVS+BR+()()()(c)()()()(d)表達式推導過程Im2col變換26圖算融合的搜索空間帶來了大量新的變換機會圖算融合的搜索空間帶來了大量新的變換機會清華大學 Tsinghua University新優化空間的挑戰探索圖

19、算融合優化空間面臨的挑戰1.自動算子生成開銷高利用TVM在GPU上生成典型矩陣乘算子需要30分鐘無法對所有表達式都進行自動算子生成2.圖算融合搜索空間大1.預定義算子數量較少，但存在無窮多無窮多張量表達式EinNet搜索空間搜索空間27清華大學 Tsinghua University高效探索圖算融合空間 算子庫映射將計算密集部分映射至算子庫，訪存密集部分映射至算子生成框架同時利用算子庫的高效性和算子生成框架的靈活性基于循環變量映射表，實現表達式至算子庫算子的匹配根據預定義算子的張量表達式，自動生成相關信息循環變量映射表28清華大學 Tsinghua University高效探索圖算融合空間 搜

20、索方法兩階段搜索算法探索式推導階段探索式推導階段：遍歷所有推導規則，探索有效的變化方案利用搜索深度限制搜索范圍收斂式推導階段收斂式推導階段：選擇性執行推導規則，減小搜索開銷比較當前表達式和目標算子（如Matmul）表達式，計算表達式距離表達式距離執行使距離減小的推導規則，加速實現到目標算子的匹配29清華大學 Tsinghua University基于表達式的圖算整合優化框架整體設計計算圖表達式自動推導模塊代碼生成模塊基于代數運算的推導規則表達式合并與拆分表達式計算轉換算子庫匹配基于表達式的代碼生成輸入程序優化后的程序基于表達式的張量程序-翻譯代數運算搜索計算融合搜索計算密集型訪存密集型人工智能

21、芯片人工智能芯片Caffe現有編程框架現有編程框架通過代數表達式建立了圖層與算子層的橋梁，為圖算融合圖算融合的優化技術提供了有力支撐30清華大學 Tsinghua University典型優化案例1 im2col復制輸入張量，將卷積計算變換為矩陣乘計算31清華大學 Tsinghua University典型優化案例2 im2col 鏡像優化累加中間結果，將卷積計算變換為矩陣乘計算32清華大學 Tsinghua University典型優化案例3 從Conv5x5到Conv3x3增加冗余計算，將5x5卷積計算變為6x6卷積，并最終變為3x3卷積33清華大學 Tsinghua University

22、實驗結果 模型加速比在batch size分別為1和16的情況下，在常用的DNN模型和A100 GPU上可取得最高最高2.32.3倍倍的加速比34清華大學 Tsinghua University實驗結果 不同后端加速效果在專家手寫的算子庫和自動算子生成框架上均能取得顯著加速算子庫：cuBLAS、cuDNN自動算子生成框架：AutoTVM、Ansor35清華大學 Tsinghua University基于表達式推導的圖算融合優化技術通過代數表達式建立了圖層與算子層的橋梁，為圖算融合的優化技術提供了有力支撐，顯著擴大了優化空間提出了基于代數表達式推導的張量程序自動優化技術，可以發掘出未被現有優化工作找到的優化策略提出了面向表達式的代碼生成方法，同時保證優化的高效性與靈活性將計算密集型的部分映射到現有算子庫以保證其高效性將訪存密集型的部分通過其表達式進行自動代碼生成來保證其功能上的靈活性36清華大學 Tsinghua University謝謝！PACMAN