程式分析與優化 - 9 附錄 XLA的緩衝區指派

周榮華發表於2022-07-02
優化

本章是系列文章的案例學習，不屬於正篇，主要介紹了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 記憶體分配統計

 

相關文章