PlaidML 中 Kernel 自動產生機制心得

來源: http://vertex.ai/blog/kernel-generation

這篇原文提供了 PlaidML 是如何將 frontend (像是 Keras 或 ONNX)中的一個 Graph 所需要的 operation, 轉換為一個優化後的 OpenCL kernel 其過程的概要. 由於關於 OpenCL code 的自動化生成, 因此個人很有興趣, 因此將本文的研讀過程做了重點的意譯.

多數的 kernel 產生演算的步驟都在於優化 caching. 主要導入了下列方法:

• 展平(flatten) Tile 程式碼(這裡的 Tile 並非影像處理使用的 tiling, 而是 PlaidML 用以實作 operation 的語言), 並將多維度 “tensor indices” 轉為 “pointer offset”
• 將龐大的 tensors 分割為能夠完整 cached 的 tiles, 並且針對使用的特定硬體選擇 tile size.
• 此外調整 tile layout 來減少 cache 中的 bank conflick 並提高符合 MAC 的資料規範.
• 以多執行緒方式執行 kernel 來利用 GPU 的 SIMD 處理器.
• 最後建立 kernel 的 Semtree 表示結構, 這將用來做 codegen. ( Semtree or Semantic Tree 是一個中介表示,以其可執行語法來描述每個函數. 其抽象層及介於 AST 與低階 LLVM-IR.)

文中以 3x3 Convolution 並接連著相同 Padding 的 ReLU 操作為例. 作者認為這樣的操作夠簡單具體, 此外也足夠複雜賴展示 PlaidML 如何解決 kernel 生成的挑戰以及如何解決常見的複雜問題, padding 與將作多個運算合併成為單一 kernel. 此外這個直接的例子其效能通常關鍵地影響整個網路的效能.

首先是以 PlaidML 的 Tile 撰寫 Convolution + Relu 的 code 如下:

function (D[N, X, Y, CI], K[I, J, CO, CI]) -> (R) {
  O[n, x, y, co : N, X, Y, CO] = +(D[n, x+i-1, y+j-1, ci] * K[i, j, co, ci]);
  R = (O > 0 ? O : 0);
}

接著在 kernel phase compilation 之前, 需要一些前置的轉換. 文中有簡短地說明:

• 自 frontend (Keras/ONNX)的 network graph 已被翻譯轉為 Tile 程式碼.
• Tile 程式碼經過分析並且決定哪些 operations 被合併入一個單一 kernel. 這通常被稱為 kernel fusion. Element-wise 的操作會被以 post-ops 方式來加到較複雜的 convolution 主要流程.
• 相關的 tensor 大小己經被決定, 並且所有的常數(strides, # of batch elements, etc) 已經是確定的.
• kernel 已經藉由 reduce/defract 轉換為固定的形式. (他們有另一篇文章介紹說明)
• 所有的 index 變數範圍都已經確定, 並且任何 index 上的額外非由範圍所表示的的限制也都已被收集.

除了範例的 Tile 程式碼外, 需要挑出所有具體相關的數字. PlaidML 其中一個有點特別的原理是它基於 input tensor 的 shape 來產生客製的 kernels. 這讓優化有著較好的效能. 像是, batch size 1 的 convolutions 可能為了效率需要一個相較於 batch size 32 完全不同的 tiling & loop 次序. 明顯地 3x3 convolution 需要的妥協不同於 7x7 convolution.

因此, 若 batch size 為 32, 輸入與輸出的 channels 數目為 64, 以及一個 224x224 的影像, 當在 kernel generation phase 的時候, 已經確定了:

• 將合併 ReLU 到 convolution kernel 之中
• 不同的 index 的範圍:

Index     Range
   ci        64
   co        64
    i         3
    j         3
    n        32
    x       224
    y       224

• 有著額外無法自 index 範圍取得的限制:

0 <= x+i-1 < 224
0 <= y+j-1 < 224

展平 (Flattening)

下一個 kernel generation 的步驟為展平(flattening). 展平將 Tile 程式碼轉換為一個簡單的表格. 基本上, 提供全部關於每個 tensor 其 memory layout 的資訊, 如此能將 tensor 的存取轉為一個由 index 組成的明確多項式所表示的記憶體位置. 這個想法是 PlaidML 優化概念的起點, 以此後續再分析這些 index. 一個簡單的例子是, 想像一個 16x16 矩陣 M, 在記憶體中以 packed C-style buffer 所表示著. 若需要存取 M[i, i+j], 能夠只存取 M 中位於記憶體 16*i + 1*(i + j) == 17*i + j的單元. 若每個 tensor 維度的存取是一個 index 組成的明確多項式, 那麼所有的記憶體存取僅只是一個明確多項式的計算. 因此被展平, 整個 convolution 就只是一個數字的集合: 對每個 index 變數, 以及每個 tensor, 何者為 index 的乘數(或是 numpy 中的 stride )? 在這範例, 再次假設為標準 C-style layout, 如此得到:

Contraction kernel_ID_0:
O[n, x, y, co : N, X, Y, CO] = +(D[n, -1 + i + x, -1 + j + y, ci] * K[i, j, co, ci])
             Range         O         D         K
      ci        64         0         1         1
      co        64         1         0        64
       i         3         0     14336     12288
       j         3         0        64      4096
       n        32   3211264   3211264         0
       x       224     14336     14336         0
       y       224        64        64         0
     off                   0    -14400         0

這裡 off 表示固定的偏移(offset) (因為在此例中藉由常數 -1 用來置中convolution).

挑出 tile size

文章段落點出 GPU 的使用需要注意的事情, 大量連續的讀取, 使用 local memory, 最後是避免 bank conflick. 而這些所有的關鍵在於 group operation 的大小, 如何儘可能重複使用自 global memory 讀進的資料. 為了有組織地處理此問題, 首先必須考量所有 index 的合理範圍. 以此例來說, 一次的 convolution 一共有 64×64×3×3×32×224×224 = 59,190,018,048 multiply-accumulates 運算. 想像將這些 operations 分群為多個 index 在各自區域範圍的 tiles. 對特定的 tile size 能夠確定:

• 需要自每個 input 讀取多少數值, 以及需要寫出多少數值到每個 output
• 這些數值會佔去多少 cache (必須考量 layout 要避免 bank conflict)
• 如何能連續的讀取/寫入 (DRAM 連續的讀寫較有效率)
• 有著多少的計算密度, 像是 flops per tile?

PlaidML 考量了大量的 tile sizes, 並對於每個 tile size, 評估這些數值. 接著使用一個所針對特定硬體的 model 來計算 tile size 的 cost. 硬體 model 包含了如下的資訊:

• L1 cache size
• 最小長度能獲得最大效能的 memory read. (基本上是 cache line 寬度)
• 相對的compute/memory 頻寬
• cache 是否為 programmer 所管理 (像是 GPUs) 或是 LRU (通常是 CPU)?

回到這個例子上, 可以看到 size optimizer 提供的部份 log:

Computing cost for tile size: [8, 32, 1, 3, 16, 8, 2]
  Compute score: to=236760072192 wg=12544 il=24 sm=20656 or=32768 mr=19456 mw=32768 op=256 rp=0 tu=256
  over memory
Computing cost for tile size: [8, 32, 1, 3, 16, 4, 4]
  Compute score: to=236760072192 wg=12544 il=24 sm=16272 or=32768 mr=15360 mw=32768 op=256 rp=0 tu=256
  over regs
Computing cost for tile size: [16, 32, 1, 1, 16, 2, 2]
  Compute score: to=236760072192 wg=50176 il=36 sm=6488 or=8192 mr=6144 mw=8192 op=256 rp=0 tu=256
  flops_per_byte=20.5714 occupancy=20
  roof_ratio=1 occ_ratio=1 thread_ratio=1 score=1
Computing cost for tile size: [8, 32, 2, 1, 16, 2, 2]
  Compute score: to=236760072192 wg=50176 il=48 sm=5264 or=8192 mr=5120 mw=8192 op=256 rp=0 tu=256
  flops_per_byte=18.5806 occupancy=20
  roof_ratio=0.929032 occ_ratio=1 thread_ratio=1 score=0.929032

對於範例所選擇的 tile size 為: [ci: 8, co: 32, i: 2, j: 3, n: 16, x: 2, y: 2]

Memory Layout

基於所挑選出的 tile size, 接著就是將工作切割為 work groups, 而每個 group 都會產生一部份的輸出結果. 在執行時每個 work group 都會得到自 GPU 傳入對應的座標, 以此指定需要計算的輸出範圍.

處理的原則上每個 inner loops 直接負責 multiply accumulations, 然後以多執行緒的方式平行處理. 然而資料如何載入 L1 以及 inner loop 如何處理必須先確立. 這裡先以 L1 cache 如何規劃開始. 這部份做了兩件事:

• 確定對於每個 inner loop 需要讀入多少資料, 以及如何載入.
• 當有兩個 index 變數以相對的方式讀取記憶體, 在載入資料時必須一併考慮兩者 — merging indexes

這樣的步驟會個別套用在 convolution 的 input(D) 與 kernel(K) 上並得到兩個表格, 對於 input 得到:

Range    Global     Local
 i_x         3     14336         1
 j_y         4        64       393
  ci         8         1        49
   n        16   3211264         3

而對於 kernel 得到:

Range    Global     Local  
 j_i         6      4096       265 
  ci         8         1        33
  co        32        64         1

在這個例子中, 輸出大小為 2×2×16×32 = 2048 單位. 由於使用了 256 threads (一個通常對於 GPU 來說好的數字, 取決於 hardward model), 每個 thread 會使用 8 個 accumulators. 這個 tile size 被選擇的原因在於:

1. 所需的 shared/local memory 需求數量小於可獲得的大小
2. 每個 thread 使用的 accumulators 數量夠小而不至於超過 register 的數目.
3. 記憶體讀取的寬度足夠 (在此例為 8 單位) 以達到對於硬體最佳的頻寬.
4. 對於每一記憶體載入而言, 計算 accumulation (compute) 的數目夠大.

這一步驟可以說是最關鍵的一步, 因為後續的 Threading 與 Codegen 只是依照此結果去做對應的程式碼輸出. 而緊接著談的是 thread layout

Threading

到此後, work group 中的 thread 有著一些需要考量的流程形式

Clear output accumulators  // Threaded
For each outer loop index:
  Load input 1 (Image)  // Threaded
  Load input 2 (Kernel)  // Threaded
  Sync()  // Make sure data is loaded before computing
  Accumulate inputs into the output (main compute)  // Threaded
  Sync()  // Make sure computations are done before loading new data
Perform post-accumulation operations and write outputs  // Threaded

對於上述結構的每一步的考量就相當貼近實務上 kernel 撰寫的優化. 主要考量為提高 GPU utilization, 減少任意兩 thread 之間的 synchronization. 連續地讀取 global memory, 避免 workgroup 中 divergence / branch 的情況.

在這裡使用了一個技巧, 建立 logical order, 其實也就是藉由 work group 中的 thread id 來交錯一些工作的執行. 類似這樣的作法也能套用在 bank conflick 的處理.

Codegen
由於至此即處理的關鍵的問題, 因此最後一個 Codegen 就不在此詳細列出. 有興趣者可以閱讀原文, 了解 codegen 中的數個步驟, 文中甚至列出了所產生的尚未優化的 CL kernel.

而最後是一些個人意見, 儘管談到了 SIMD, 但是 PlaidML 似乎單純認為 thread/lane 是對於 SIMD processor 固定(良好的)考量方式, 並沒有特別考量使用 vector. 觀察了其 opencl backend, 似乎也只在 fp16 的時候有考量使用 vector.