本章是系列文章的案例學習，不屬于正篇，主要介紹了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 記憶體分配統計
轉載請註明出處，本文鏈接：https://www.uj5u.com/ruanti/498993.html
標籤：其他
下一篇：嵌入式軟體架構設計-函式呼叫