DeepRoute Lab | CUDA算子优化：转置篇

转置是深度模型中应用最广泛的算子之一。各种深度学习框架崛起的今天，它被“隔离”到了框架的底层，不再显式的活跃于用户面前。谈及转置，就不得不先聊聊卷积的优化算法以及它们所需要的图像数据格式。

AUTHOR：vector  小何

01

图像数据格式的前世今生

目前业界对于卷积的实现主要有4种方式，大白算法(Direct Convolution)，Winograd Convolution，隐式矩阵乘卷积(Implicit GEMM Convolution)和快速傅里叶变换(FFT Convolution)，前两个需要的数据格式是通道前置(NCHW)，而后两者需要通道后置(NHWC or NCxHWx)。

1.1 Direct Convolution vs Implicit GEMM Convolution

大白算法，顾名思义，就是用最直观的方式对原始数据(Feature Map)和卷积核(Filter)做互相关(Cross Correlation)计算。这种算法的本质逻辑就是卷积核在原始数据上不断地滑动计算出结果。从内存加载数据到寄存器的过程中，为了保证访存效率，需要H, W维度连续，因此对于大白算法来说，NCHW数据格式显然占优。

隐式矩阵乘卷积，还有另外一个耳熟能详的名字 ------ im2col。既然卷积实质上就是输入通道维度上的累加，那么不妨以此作为优化突破口，把累加维度后置，方便SIMD指令访存计算。对于此种优化算法来说，NHWC数据格式更占优。

1.2 CPU vs GPU

对于Intel系列的CPU来说，得益于SIMD指令集SSE系列，AVX系列的加持，隐式矩阵乘在性能上全方位碾压了大白算法。因此许多CPU线性代数计算库会以NHWC作为标准数据格式。

以Nvidia为代表的GPU经常被用于各种高性能计算场景，在volta架构之前，N卡的计算核心是未分家的ALU和FPU，习惯性会将其统称为CUDA Core。既然仍然依赖于常规FPU的计算，那么就脱离不开乘法指令比加法指令执行更慢的魔咒，因此Winograd卷积成为了最高效的卷积算法。如果对于精度要求较高或输入输出形状极其特殊的场景来说，高并行度的大白算法是妥协后的选择。总的来说隐式矩阵乘仍然排不上号。

从Volta架构开始，N卡引入了新的硬件结构：张量计算单元(1st generation Tensor Core)，专门用于加速FMA-reduction(Fused Multiply Add with Reduction)类型的算子。同时隔壁Google家也早早地就用上了自研的TPU(Tensor Process Unit)来加速深度模型的训练和推理。这些主流的张量计算单元是通过脉动阵列(Systolic Array)技术实现的，这项古老的技术在沉寂了多年之后，直到急需高效FMA-reduction算子的今天再次得以重见天日，随之而来的是隐式矩阵乘算法的快速崛起。自然而然的，通道后置格式也变得重要了起来。

几种数据格式间的高效转换 ------ 转置变得非常重要。

1.3 转置

让我们先来简化一下问题，实际上(N, C, H, W)也可以写作(N, C/x, x, H, W)，而我们需要的数据排布为(N, C/x, H, W, x)。

这个时候我们把没有发生变化的维度(N和C/x, H和W)合并。

我们就将问题简化成了：

我们不难发现，这个问题的本质实际上就是多个普通的二维矩阵转置。在NCHW转置到NC/xHWx时，针对于不同的数据类型，编译期常量x的值是固定的，常用的：

因此，对于int8类型的数据来说，我们待解决的问题就变成了：

在这里，我们主要以针对于int8数据类型的NCHW转置NC/32HW32为切入口，来浅谈CUDA优化。

02

初版实现

截止到CUDA 11.8 和 Hopper架构之前， CUDA在软件层面分为4个层级网格(Grid)，线程块(ThreadBlock)，线程束(Warp)，线程(Thread)。每个流处理器有4个线程束调度器(Warp Scheduler)，因此每个线程块分配128线程一般是最高效的分配方案。接下来，定好每个层级需要处理哪部分数据即可。在这个问题中，我们自然而然会想到每个Block处理32行数据，每个线程负责搬运一个数，那么图示如下：

（点击查看大图）

接下来，我们只需要计算出每个线程需要从何取数又放回何处即可。先来定义描述各层级Shape的结构体：

（点击查看大图）