本章是系列文章的案例學習，不屬於正篇，主要介紹了TensorFlow引入的XLA的優化演算法。XLA也有很多局限性，XLA更多的是進行合併，但有時候如果參數特別多的場景下，也需要進行分割。XLA沒有數據切分的功能。當前最主流的編譯器領域的編譯優化功能還是mlir。 本文中的所有內容來自學習DCC888 ...
本章是系列文章的案例學習，不屬於正篇，主要介紹了TensorFlow引入的XLA的優化演算法。XLA也有很多局限性，XLA更多的是進行合併，但有時候如果參數特別多的場景下，也需要進行分割。XLA沒有數據切分的功能。當前最主流的編譯器領域的編譯優化功能還是mlir。
本文中的所有內容來自學習DCC888的學習筆記或者自己理解的整理，如需轉載請註明出處。周榮華@燧原科技
9.1 什麼是XLA
- XLA是Accelerated Linear Algebra的簡稱。
第一次看到Accelerated被簡稱為X的時候，有點奇怪，因為Accelerated裡面可沒有一個字母是X，但Accelerated的發音和X相同，這樣簡化之後可以避免一個簡寫中存在多個A的不協調，XLA讀起來確實比ALA朗朗上口一點:)
- XLA - TensorFlow, compiled. Mar. 6th, 2017
- XLA是一種編譯線性代數領域相關的編譯器，主要用來加速TensorFlow的模型優化和目標代碼生成
- 除了TensorFlow外，XLA也可以用在多種前端中，包括TensorFlow，Pytorch, JAX, Julia和Nx
- XLA的功能設計上其實與target arch無關的，所以也可以支持多種後端：CPU，GPU或者其他硬體
在2017年XLA誕生的時候，那時給出的幀處理加速數據如下：
帶來相應加速效果的主要因素是通過分析和調度記憶體使用，刪除了一些中間表達的存儲緩存，其中一個主要的方法就是緩衝區指派演算法，也就是本文主要準備描述的。
XLA的設計理念是一種近似SSA的中間表達：
- 變數只能被初始化（除了初始化，不能額外修改）
- 更短的生命周期
- 清晰的Def-Use鏈
XLA: Optimizing Compiler for Machine Learning | TensorFlow中有個油管視頻詳細講解了XLA的原理，通過這個也可以理解一下TensorFlow的原理：
tf.function → tf2xla橋 → 優化前的xla hlo → xla的一些列優化 → 優化後的hlo → 可執行binary → tf2xla橋 → tf runtime → target arch上執行
9.2 靜態記憶體分配分析
9.2.1 為什麼可以做分析
- 靜態計算圖本身的特性
- 張量在執行階段只會使用固定的記憶體空間
- 靜態計算圖在執行前就可以靜態推斷
9.2.2 靜態記憶體分析的優勢
- 為運算元提供通用的記憶體分配
- 重用前面運算元的記憶體，減少重新分配和拷貝過程
- 減少額外的碎片和記憶體管理
9.2.3 靜態記憶體分析的局限性
- 僅針對靜態計算圖有效
9.3 緩衝區管理的目標
- 儘可能重用記憶體
- 當記憶體不足以完成任務時報錯
緩衝區定義：每個運算元定義一個緩衝區
緩衝區申請、支配原則：
- 在生命周期上不相互干擾的緩衝區可以使用相互覆蓋的記憶體
- 如果緩衝區和其他記憶體都衝突，需要重新申請記憶體並指派給它
- 所有申請的記憶體按組存放
9.4 緩衝區分析（有可以稱為別名分析）
緩衝區分析的過程和指針分析的過程有很多類似的地方，所以很多地方又稱為別名分析。
一個IR需要定義≥1個邏輯緩衝區
用{def, {}} 來定義一個緩衝區
緩衝區{b, {}} 和 {b, {1}}可以相互覆蓋
來自不同IR的邏輯緩衝區可以復用同一塊記憶體
例如對下麵的偽代碼，可以知道d和b是別名關係，因為它們指向同一片記憶體：
9.4.1 定義所有指令的所有邏輯緩衝區
按拓撲順序遍歷（選擇什麼順序？逆後根排序）計算圖，為每個指令分配緩衝區，例如上面的偽代碼，生成緩衝區如下：
1 Buffer(a, {}) : [ (a, {}) ] 2 Buffer(b, {}) : [ (b, {}) ] 3 Buffer(c, {}) : [ (c, {})]
9.4.2 HLO內部的別名分析後的結果
1 Buffer(a, {}) : [ (c, {0}), (a, {}) ] 2 Buffer(b, {}) : [ (c, {1}), (b, {}) , (d, {}) ] 3 Buffer(c, {}) : [ (c, {})]
9.4.3 跨HLO的別名分析
基於近似SSA的HLO語法定義，編譯過程變得簡單了很多（SSA化是很多編譯中的主要工作）
9.4.4 基於上面偽代碼的生命周期干擾分析
從下麵生成的圖來看，a和b互相干擾，不能公用緩衝區，e理論上是d的拷貝別名，所以和b也是別名關係。和寄存器分配不同的是，考慮到多線程執行場景，不同流中要用到的緩衝區不能分配到同一個組，所以a/b雖然和e在下麵的計算圖中沒有干擾，但由於e是後面HLO的輸入，所有e不能和當前計算圖中的任意一個緩衝形成別名關係。
9.5 buffer指派的功能
9.5.1 將能夠重用的buffer儘可能重用
沒有生命期干擾的緩衝區都可以分配到同樣的記憶體
9.5.2 緩衝區分配複合
將著色相同的緩衝區複合到一起申請（可能不同緩衝區占用某個實際緩衝區的不同部分，但大家相互之間的關係決定了它們可以相鄰申請）
9.5.3 從全局分析去掉記憶體碎片
9.5.4 峰值記憶體壓力預測
9.5.5 記憶體分配統計
