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

來源丨 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_）：

簡單來說就是我們需要根據(jù)index的索引完成對當(dāng)前Tensor dim維度的inplace加和，注意被加數(shù)是由另外一個Tensor src決定的。在PyTorch的codebase中搜索index_add，我們可以發(fā)現(xiàn)這個操作應(yīng)用得非常廣泛，比如說作為as_strided算子的backward的一部分，作為一些sparse操作的一部分等等。我最近研究了一下，發(fā)現(xiàn)PyTorch對index_add算子的cuda kernel進(jìn)行了較為精細(xì)的優(yōu)化，主要有兩個亮點，本篇文章就來學(xué)習(xí)一下。順便提一下，在PyTorch中index_add的cuda kernel實現(xiàn)在https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L712 ，如果你想自己詳細(xì)讀這個代碼我建議先編譯一下PyTorch再進(jìn)行調(diào)試和閱讀，編譯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的計算邏輯可能會因為重復(fù)訪問index導(dǎo)致全局內(nèi)存訪問過多，而如果index很大那么為了保證性能kernel又需要滿足足夠的并行度才可以。為了平衡這兩種情況，PyTorch按照index的元素個數(shù)實現(xiàn)了2套kernel。這2套kernel的實現(xiàn)在：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L576-L675 。然后根據(jù)index元素個數(shù)進(jìn)行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_）為例記錄了各個中間變量的值，并加上了一些方便理解的注釋，感興趣的可以查看。我們這里展示一下當(dāng)index的元素很少的時候的indexFuncSmallIndex kernel實現(xiàn)（代碼中的設(shè)置是index元素個數(shù)少于16）：

// 如果索引的數(shù)量很少，我們更喜歡使用這個Kernel來避免重新加載 index。
// 這個kernel實際上適用于幾乎所有問題大小的選擇，但如果選擇的索引數(shù)量很大，
// 那么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.
  // 為了避免重新加載我們正在復(fù)制的索引，加載一次以處理所有正在選擇的點，以便盡可能地重復(fù)使用它。 
  // 當(dāng)這是一個不錯的選擇（選擇的索引數(shù)量很少）時，就會選擇這個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)來避免重復(fù)加載 index Tensor（這個時候index Tensor信息由indices管理），后續(xù)的實驗結(jié)果也將證明這個優(yōu)化在 index 元素個數(shù)比較小而 self Tensor 比較大的時候是有一定性能提升的。然后選定一個indices[i] 之后就啟動一堆線程計算完這個indices[i]對應(yīng)的 self Tensor的一個切片（linearIndex < innerSize）。indexFuncLargeIndex Kernel我就不展示了，感興趣的小伙伴可以直接閱讀代碼實現(xiàn)。實現(xiàn)完這兩個Kernel之后，我們可以在 https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L753-L778 這里看到PyTorch分別為這兩個Kernel設(shè)置了不同的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以正在索引的每個切片的大小為基準(zhǔn)來設(shè)定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 張量的總大小為基準(zhǔn)來設(shè)定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的元素個數(shù)比較小也就是smallIndex的情況，線程快的數(shù)量由sliceSize來決定，而對于index元素個數(shù)比較大也就是largeIndex的時候線程塊的數(shù)量則由輸入Tensor self的總元素數(shù)量來決定。我個人感覺這里設(shè)置GridSize和BlockSize還是存在一定問題的，在profile的時候ncu對于index比較小并且輸入Tensor也不太大的情況會提示grid太小無法充分發(fā)揮并行性的問題。建議閱讀https://mp.weixin.qq.com/s/1_ao9xM6Qk3JaavptChXew 這篇文章設(shè)置更合理的BlocSize和GridSize，或許可以提升smallIndex Kernel的性能。比如index很小但是輸入Tensor只有一個維度的情況下，這個時候PyTorch只會啟動一個Block以及一個Thread，這顯然是個bad case：在這里插入圖片描述

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

我這里對比了一下PyTorch的index_add和oneflow中index_add的性能表現(xiàn)。做性能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，測試結(jié)果如下：

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

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

• 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