程式分析與優化

本章是系列文章的案例學習，不屬於正篇，主要介紹了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 緩衝區分配複合

將著色相同的緩衝區複合到一起申請（可能不同緩衝區佔用某個實際緩衝區的不同部分，但大家相互之間的關係決定了它們可以相鄰申請）