如何实现比 PyTorch 快 6 倍的 Permute/Transpose 算子？

无论是在统治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