中文字幕久久久人伦,玖玖婷婷,午夜精品导航,手机在线观看av网站,www.国产免费,免费日本在线

無論是在統治nlp屆的transformer，還是最近視覺領域的新秀vision transformer，我們都能在模型中看到transpose/permute算子的身影，特別是在多頭注意力機制(multi-head attention)中，需要該算子來改變數據維度排布。

顯然，作為一個被高頻使用的算子，其cuda實現會影響到實際網絡的訓練速度。本文會介紹優化permute kernel的技巧，并跟pytorch的permute，原生的copy操作進行實驗對比。
1樸素的permute實現
permute算子的作用是變換張量數據維度的順序，舉個例子：
x = flow.randn(2, 3)
y = x.permute(1, 0)
y.shape
(3, 2)

其實現原理也可以很容易理解，即輸出tensor的第i維，對應輸入tensor的dims[i]維，上述例子中 permute 實現對應的偽代碼如下：
for row in x.shape[0]:
for col in x.shape[1]:
y[row][col] = x[col][row]

但是實際情況與上面的偽代碼有出入，張量的shape是數學上的概念，在物理設備上并不真實存在。

張量的數據都是保存在一塊連續的內存中，下圖分別從上層視角和底層視角描述了形狀為(2, 3)的張量的存儲方式：

permute實現原理為：

• 通過當前輸出的一維偏移量(offset)計算對應的高維索引

  
  

• 然后根據參數dims重新排列輸出索引，進而得到輸入索引。

  
  

• 將輸入索引轉換成輸入偏移量

  
  

• 最后進行數據移動，整個過程的示意圖如下：

  
  

完成permute后，輸出如下圖所示：

整個 permute 計算過程需要經過多次一維偏移量offset和高維索引之間的轉換，為了避免一次次手工計算，提供了一個工具類ndindexoffsethelper來方便做上述轉換。

2ndindexoffsethelper

ndindexoffsethelper的主體方法如下：

• ndindextooffset方法把高維索引轉為一維偏移量

  
  

• offsettondindex方法把一維偏移量轉為高維索引

  
  

有了這么一個工具類，那我們就可以很輕松的寫出一版naive permute kernel了，核函數如下：template
__global__ void permutekernel(permutekernelparams params) {
using t = typename std::aligned_storage::type;
const t* src = reinterpret_cast(params.src);
t* dst = reinterpret_cast(params.dst);
indextype src_index[num_dims];
indextype dst_index[num_dims];
cuda_1d_kernel_loop_t(indextype, i, params.count) {
params.dst_index_helper.offsettondindex(i, dst_index);
#pragma unroll
for (size_t dim = 0; dim  (2, 3, 0, 1)
x = flow.randn(3, 4, 5, 6)
y = x.permute(2, 3, 0, 1)
y.shape
(5, 6, 3, 4)
顯然這是一個四維的permute情形，但這里第2，3維，第0，1維是一起permute的，所以我們可以看成是一種二維的permute情形：
# (0, 1, 2, 3) -> ((2, 3), (0, 1))
x = x.reshape(x.shape[0]*x.shape[1], x.shape[2]*x.shape[3])
y = x.permute(1, 0)
y = y.reshape(x.shape[2], x.shape[3], x.shape[0], x.shape[1])

合并維度后，在利用ndindexoffsethelper根據偏移量計算索引時，合并前需要計算成四維索引，而合并后我們只需計算成二維索引。相比合并前減少除法和乘法的次數，進而提升速度。

3. 使用更大的訪問粒度

細心的朋友們可能觀察到核函數中有一個模板參數size_t movement_size，它表示的是訪問元素的粒度。
在nvidia性能優化博客increase performance with vectorized memory access中提到可以通過向量化內存操作來提高cuda kernel性能，能夠減少指令數，提高帶寬利用率。（鏈接：https://developer.nvidia.com/blog/cuda-pro-tip-increase-performance-with-vectorized-memory-access/）
我們設置訪問粒度的規則如下：

• cuda支持的訪問粒度為1b，2b，4b，8b，16b，粒度越大性能越好

  
  

• 最后一個維度是作為整體來移動的，即permutation[n-1]==x.dims[n-1]，且大小是新訪問粒度的倍數

  
  

• 保證數據指針滿足新訪問粒度的對齊要求

  
  

針對規則2，對應著以下permute場景：(0, 1, 2, 3) -> (0, 2, 1, 3)其中最后一維并沒有變化，僅僅是第1，2維進行交換，那么我們可以使用更大的訪問粒度來讀取數據，再進行permute操作。代碼中通過getmovementsize函數來確定訪問粒度的大小。
我們使用nsight compute對pytorch的permute和原生copy操作對比測試運行時間和帶寬，測試結果如下：

其中測試環境為nvidia a100 40gb，場景為(0, 1, 2)->(1, 0, 2)，橫坐標表示數據形狀及數據類型。測試數據覆蓋了16mb到128mb不同大小的數據，數據類型包含fp32和half兩種類型。
從上面兩張圖可以看到，在大部分情況下都可以逼近甚至略高于copy操作的帶寬。與pytorch對比，在操作耗時上最少快1.24倍，最快能達1.4倍。這里permute的帶寬比原生copy還高一點，是因為copy kernel里沒有做unroll指令間并行優化，而permute kernel內部做了相關優化，這里僅做參考。使用上面的兩個優化技巧，就能輕易做到比pytorch的實現要快了。常規的permute適用情況比較廣泛，也因此可能存在訪存不合并的情況。在一些特殊的場景下，我們可以通過合并訪存以提升帶寬利用率和速度，這就引出我們下個關于batchtranspose優化的話題。
4batchtranspose優化
batchtranspose操作即矩陣轉置，僅交換矩陣最后的兩維，以下情況均符合batchtranspose的定義，其中括號內容表示維度的順序：
(0, 1) -> (1, 0)
(0, 1, 2) -> (0, 2, 1)

在樸素的permute方案中，對于最后一維作為整體移動的情況下，已經進行充分的優化。但實際場景中還存在矩陣轉置的情況，此時無法應用第三條增大訪問粒度的優化操作，并且不滿足訪存合并要求，導致性能不佳。以pytorch為例，在數據大小為128mb情況下進行batchtranspose時，因為未合并的訪存導致實際讀取數據量遠大于寫入數據量(7-8倍)。

在英偉達性能優化博客an efficient matrix transpose in cuda c/c （https://developer.nvidia.com/blog/efficient-matrix-transpose-cuda-cc/）中，其做法是設置一塊shared memory，然后將一行數據讀取到shared memory，再按列順序將shared memory中的元素寫回到global memory中。得益于shared memory訪問粒度小的特性(global memory是32b，shared memory是4b)，進而避免global memory的訪存不連續的問題。
shared memory相比global memory有15倍更高的帶寬，20-40倍更低的延遲，因此額外引入的讀寫開銷可以忽略不計。
此外我們給shared memory多padding了一個元素，進而讓以列順序訪問的元素能夠均勻分布在32個bank上，避免bank conflict。對應的示意圖如下(其中灰色部分代表padding元素)：

基于上述提到的點我們實現了一版batchtranspose，代碼如下：
template
__global__ void batchtransposekernel(const void* src_ptr, void* dst_ptr, indextype h, indextype w,
indextype num_tile_rows, indextype num_tile_cols,
int32_t block_nums) {
using t = typename std::aligned_storage::type;
__shared__ t tile[tile_size][tile_size 1]; // to avoid bank conflict.

const t* src = reinterpret_cast(src_ptr);
t* dst = reinterpret_cast(dst_ptr);

indextype batch_num_tile = num_tile_rows * num_tile_cols;
for (int i = blockidx.x, step = griddim.x; i