徹底搞懂PyTorch index_add 操作涉及的優(yōu)化技術

來源丨 GiantPandaCV

本文把pytorch index_add算子的代碼抽取出來放在：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/indexing/index_add_cuda_pytorch_impl.cu 。如果不太熟悉PyTorch的話也可以直接看這個.cu文件，有問題請在這個repo提issue。

我們可以在 PyTorch 的文檔中找到 torch.index_add_ 的定義（https://pytorch.org/docs/stable/generated/torch.Tensor.index_add_.html#torch.Tensor.index_add_）：

簡單來說就是我們需要根據index的索引完成對當前Tensor dim維度的inplace加和，注意被加數是由另外一個Tensor src決定的。在PyTorch的codebase中搜索index_add，我們可以發(fā)現這個操作應用得非常廣泛，比如說作為as_strided算子的backward的一部分，作為一些sparse操作的一部分等等。我最近研究了一下，發(fā)現PyTorch對index_add算子的cuda kernel進行了較為精細的優(yōu)化，主要有兩個亮點，本篇文章就來學習一下。順便提一下，在PyTorch中index_add的cuda kernel實現在https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L712 ，如果你想自己詳細讀這個代碼我建議先編譯一下PyTorch再進行調試和閱讀，編譯PyTorch源碼可以參考：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/tree/master/how-to-complie-pytorch-from-source（這個也是參考PyTorch官方的教程，補充了幾個報錯的坑） 。

PyTorch優(yōu)化的出發(fā)點是，index_add操作中index這個Tensor是尤其重要，它決定了輸入Tensor的哪些位置會被重新賦值，然后index的元素可多可少。如果使用同一套naive的計算邏輯可能會因為重復訪問index導致全局內存訪問過多，而如果index很大那么為了保證性能kernel又需要滿足足夠的并行度才可以。為了平衡這兩種情況，PyTorch按照index的元素個數實現了2套kernel。這2套kernel的實現在：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L576-L675 。然后根據index元素個數進行dispatch：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L801-L829 。在這里插入圖片描述我在 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/indexing/indexing_pytorch_explain.cu#L381-L505 這里以PyTorch文檔展示的例子（https://pytorch.org/docs/stable/generated/torch.Tensor.index_add_.html#torch.Tensor.index_add_）為例記錄了各個中間變量的值，并加上了一些方便理解的注釋，感興趣的可以查看。我們這里展示一下當index的元素很少的時候的indexFuncSmallIndex kernel實現（代碼中的設置是index元素個數少于16）：

// 如果索引的數量很少，我們更喜歡使用這個Kernel來避免重新加載 index。
// 這個kernel實際上適用于幾乎所有問題大小的選擇，但如果選擇的索引數量很大，
// 那么indexFuncLargeIndex Kernel是增加并行度的更好選擇。
// 下面的innerSize就是輸人的self張量忽略dim維度的切片大小，對于每一個indices[i]，我們都要處理innerSize個元素的copy

// selfAddDim(dstAddDim) = 0
// sourceAddDim(srcAddDim) = 0
// sliceSize(innerSize) = 3
// selfAddDimSize(dstAddDimSize) = 5
// selfNumel(dstNumel) = 15
// selfInfo.sizes(dst): 1, 3, 
// selfInfo.strides(dst): 3, 1,
// sourceInfo.sizes(src): 1, 3, 
// sourceInfo.strides(src): 3, 1
// indexInfo.sizes(indices): 3, 
// indexInfo.strides(indices): 1,

template <typename T, typename IndicesType, typename IndexType, int DstDim, int SrcDim, int IdxDim,
          typename func_t>
__global__ void indexFuncSmallIndex(cuda::detail::TensorInfo<T, IndexType> dst,
                                    cuda::detail::TensorInfo<T, IndexType> src,
                                    cuda::detail::TensorInfo<IndicesType, IndexType> indices,
                                    int dstAddDim,
                                    int srcAddDim,
                                    IndexType innerSize,
                                    int64_t dstAddDimSize,
                                    int64_t dstNumel,
                                    const func_t& op,
                                    T alpha) {
  // In order to avoid reloading the index that we are copying, load
  // it once to handle all of the points that are being selected, so
  // it can be reused as much as possible. This kernel is chosen when
  // this is a good choice (small number of chosen indices), since
  // re-accessing indices in addition to src elements can be slow.
  // 為了避免重新加載我們正在復制的索引，加載一次以處理所有正在選擇的點，以便盡可能地重復使用它。 
  // 當這是一個不錯的選擇（選擇的索引數量很少）時，就會選擇這個Kernel，
  // 因為除了 src 元素之外，重新訪問索引可能很慢。
  for (IndexType srcIndex = 0; srcIndex < indices.sizes[0]; ++srcIndex) {
    // Lua indices begin at 1
    IndexType dstIndex =
        indices.data[cuda::detail::IndexToOffset<IndicesType, IndexType, IdxDim>::get(srcIndex, indices)];
    CUDA_KERNEL_ASSERT(dstIndex < dstAddDimSize);

    // We stride over the output ignoring the indexed dimension
    // (innerSize), whose offset calculation is handled differently
    for (IndexType linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
         linearIndex < innerSize;
         linearIndex += gridDim.x * blockDim.x) {
      IndexType dstOffset =
          cuda::detail::IndexToOffset<T, IndexType, DstDim>::get(linearIndex, dst);
      dstOffset += dstIndex * dst.strides[dstAddDim];

      IndexType srcOffset =
          cuda::detail::IndexToOffset<T, IndexType, SrcDim>::get(linearIndex, src);
      srcOffset += srcIndex * src.strides[srcAddDim];

      T val = src.data[srcOffset] * alpha;
      op(dst.data, dstOffset, dstNumel, &val);
    }

  }
}

我們可以看到首先有一個for (IndexType srcIndex = 0; srcIndex < indices.sizes[0]; ++srcIndex) 的循環(huán)來避免重復加載 index Tensor（這個時候index Tensor信息由indices管理），后續(xù)的實驗結果也將證明這個優(yōu)化在 index 元素個數比較小而 self Tensor 比較大的時候是有一定性能提升的。然后選定一個indices[i] 之后就啟動一堆線程計算完這個indices[i]對應的 self Tensor的一個切片（linearIndex < innerSize）。indexFuncLargeIndex Kernel我就不展示了，感興趣的小伙伴可以直接閱讀代碼實現。實現完這兩個Kernel之后，我們可以在 https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L753-L778 這里看到PyTorch分別為這兩個Kernel設置了不同的GridSize和BlockSize。

// selfAddDim = 0
// sourceAddDim = 0
// sliceSize = 3
// selfAddDimSize = 5
// selfNumel = 15

#define SMALL_INDEX(TENSOR_TYPE, INDICES_TYPE, TYPE, SELF_DIM, SOURCE_DIM, IDX_DIM)     \
  indexFuncSmallIndex<TENSOR_TYPE, INDICES_TYPE, TYPE, SELF_DIM, SOURCE_DIM, IDX_DIM>   \
    <<<smallIndexGrid, smallIndexBlock, 0, stream>>>(                                   \
      selfInfo, sourceInfo, indexInfo,                                                  \
      selfAddDim, sourceAddDim, sliceSize, selfAddDimSize,                              \
      selfNumel, reduce_add, alpha_value);                                              \
  C10_CUDA_KERNEL_LAUNCH_CHECK();

#define LARGE_INDEX(TENSOR_TYPE, INDICES_TYPE, TYPE,                        \
                    SELF_DIM, SOURCE_DIM, IDX_DIM, IDX_IS_MAJOR)            \
  indexFuncLargeIndex<TENSOR_TYPE, INDICES_TYPE, TYPE,                      \
                      SELF_DIM, SOURCE_DIM, IDX_DIM, IDX_IS_MAJOR>          \
    <<<largeIndexGrid, largeIndexBlock, 0, stream>>>(                       \
      selfInfo, sourceInfo, indexInfo,                                      \
      selfAddDim, sourceAddDim, sourceTotalSize,                            \
      (IDX_IS_MAJOR) ? sliceSize : numIndex,                                \
      selfAddDimSize, selfNumel, reduce_add, alpha_value);                  \
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  // small index以正在索引的每個切片的大小為基準來設定GridSize和BlockSize，同時要考慮到需要滿足足夠多的wave保證利用率
  const dim3 smallIndexGrid(std::min(ceil_div(sliceSize, (ptrdiff_t)128), (ptrdiff_t)(mpc * 8)));
  const dim3 smallIndexBlock(std::min(sliceSize, (ptrdiff_t)128));

  // large index以source 張量的總大小為基準來設定GridSize和BlockSize，同時要考慮到需要滿足足夠多的wave保證利用率
  const dim3 largeIndexGrid(std::min(ceil_div(sourceTotalSize, (ptrdiff_t)128), (ptrdiff_t)(mpc * 8)));
  const dim3 largeIndexBlock(std::min(sourceTotalSize, (ptrdiff_t)128));

對于index的元素個數比較小也就是smallIndex的情況，線程快的數量由sliceSize來決定，而對于index元素個數比較大也就是largeIndex的時候線程塊的數量則由輸入Tensor self的總元素數量來決定。我個人感覺這里設置GridSize和BlockSize還是存在一定問題的，在profile的時候ncu對于index比較小并且輸入Tensor也不太大的情況會提示grid太小無法充分發(fā)揮并行性的問題。建議閱讀https://mp.weixin.qq.com/s/1_ao9xM6Qk3JaavptChXew 這篇文章設置更合理的BlocSize和GridSize，或許可以提升smallIndex Kernel的性能。比如index很小但是輸入Tensor只有一個維度的情況下，這個時候PyTorch只會啟動一個Block以及一個Thread，這顯然是個bad case：在這里插入圖片描述

index_add里面的第二個優(yōu)化亮點是對Tensor的維度壓縮，對應代碼的https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L793， https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L787 ，這個維度壓縮是什么意思呢？假設index_add操作的輸入Tensor是三個維度假設形狀為(32, 1024, 1024)，而dim設置為0。那么在cuda Kernel中索引位置的時候是可以提前把dim后面的維度給合并起來的（這里使用TensorInfo數據結構來完成，其實本質上就是操作這個TensorInfo對象維護的Tensor的stride和size，具體可見這里的實現：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/CollapseDims.h#L22），這樣子原始的輸入Tensor的形狀就會變成（32, 1024）。這樣在indexFuncSmallIndex和indexFuncLargeIndex Kernel里面做坐標映射的時候就可以降低計算量以及降低對全局內存的訪問提升帶寬。注意，這里的維度壓縮也可以壓縮dim之前的所有維度為一個維度，這樣子最終kernel需要處理的self輸入張量維度只有1，2，3三種情況。雖然這個優(yōu)化是算法層面的優(yōu)化，但是也間接讓cuda kernel的帶寬進行了提升和計算量進行了下降。實際上這個思路也啟發(fā)了我在oneflow中實現index_add的kernel，我也是間接做了維度壓縮。以這個例子來說：

x = torch.randn(32, 1024, 1024).to("cuda")
t = torch.randn(15, 1024, 1024).to("cuda")
index = torch.randint(0, 32, (15,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

使用ncu在a100 pcie 40g上profile，我發(fā)現使用了維度壓縮優(yōu)化之后將這個cuda kernel從接近300+us的運行速度提升到了180+ us。

我這里對比了一下PyTorch的index_add和oneflow中index_add的性能表現。做性能profile的時候，我使用了以下腳本：

import torch

x = torch.randn(32*1024*1024).to("cuda")
t = torch.randn(15).to("cuda")
index = torch.randint(0, 1024, (15,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

x = torch.randn(32*1024, 1024).to("cuda")
t = torch.randn(15, 1024).to("cuda")
index = torch.randint(0, 1024, (15,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

x = torch.randn(32, 1024, 1024).to("cuda")
t = torch.randn(15, 1024, 1024).to("cuda")
index = torch.randint(0, 32, (15,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

x = torch.randn(32*1024*1024).to("cuda")
t = torch.randn(1024).to("cuda")
index = torch.randint(0, 1024, (1024,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

x = torch.randn(32*1024, 1024).to("cuda")
t = torch.randn(1024, 1024).to("cuda")
index = torch.randint(0, 1024, (1024,)).to("cuda")
x.index_add_(0, index, t)
torch.cuda.synchronize()

測試環(huán)境為 A100 PCIE 40G，測試結果如下：

整體來說 PyTorch 在 index Tensor元素很小，但Tensor很大的情況下相比于oneflow有一些性能提升，其它情況和 OneFlow 基本持平，也有一些case是慢于oneflow比如index很小但是輸入Tensor只有一個維度的情況下，這個時候PyTorch只會啟動一個Block以及一個Thread，這顯然是個bad case。OneFlow的index_add實現在 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/index_add_kernel.cu ，我們并沒有針對 index 的大小來單獨派發(fā)kernel，所以在某些case上性能暫時比PyTorch低一些，后續(xù)有需求的話可以繼續(xù)優(yōu)化下。

我這里相對粗糙的學習了一下調研PyTorch index_add這個算子的cuda實現的優(yōu)化技術。但PyTorch的這個index_add實現仍然有一些改進空間，比如IndexToOffset的實現有取模操作，這個可以改成一次乘法和減法，可以節(jié)省計算指令。然后index_add 的兩個kernel來說，GridSize和BlockSize并不是很合理，有改進空間。

• https://github.com/pytorch/pytorch
• https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/index_add_kernel.cu
• https://github.com/BBuf/how-to-optim-algorithm-in-cuda