转置是深度模型中应用最广泛的算子之一。各种深度学习框架崛起的今天，它被“隔离”到了框架的底层，不再显式的活跃于用户面前。谈及转置，就不得不先聊聊卷积的优化算法以及它们所需要的图像数据格式。
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的计算，那么就脱离不开乘法指令比加法指令执行更慢的魔咒，因此空洞卷积成为了最高效的卷积算法。如果对于精度要求较高或输入输出形状极其特殊的场景来说，高并行度的大白算法是妥协后的选择。总的来说隐式矩阵乘仍然排不上号。
从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的结构体：
（点击查看大图）
在定义计算存取数偏移量的仿函数。选择封装成仿函数是为了方便偏特化，以便把一些常量计算放到编译期进行，减少运行时的冗余计算指令：
（点击查看大图）
（点击查看大图）
根据我们最初的分块策略，初版的Kernel实现如下：
（点击查看大图）
其中 GlobalReadInitializer 和 GlobalWriteInitializer 用来获取每个线程块读与写全局内存的首地址。GlobalReader 和 GlobalWriter 则根据二维排布的行和列来计算每个线程读与写的地址。可能会有同学好奇为何全部使用uint64_t类型，这样做是为了对指针地址做运算时，减少各种clamp相关和cast相关的指令，可以最大化减少指令数量。当然对于64bit运算，都是分高低位的，所以害怕计算周期变久的同学大可放心使用。当输入形状为(1, 32, 576, 960)时，在GeForce 4090上的速度为69.05us，看起来确实对得起它的龟速。
功能实现完了，接下来我们开启性能优化之旅。
03访存连续与共享内存
3.1 访存连续
任何设备内存事务(Memory Transaction)的过程都是查找首位对齐地址(Aligned Address)，往后加载/写入连续(Coalesced)多个字节，具体多少字节由各级缓存的CacheLine决定。对于N卡来说，L2的CacheLine是32Bytes。非连续访存会导致设备无法将多次访存合并，导致多次内存事务。从Global到L2的指令延迟在500cycle左右的数量级，多次内存事务会导致LDST单元空闲率变高，反映出来的就是我们常说的访存效率(Memory Efficiency)变低。
不难发现，初版实现中写回数据时是连续访存，而加载数据时是非连续的。这是转置类算子的通用问题：无论如何变通，总是会在加载或写回过程中有一个过程是非连续访存。这时，我们需要换种思路来解决这个问题，如果我们让那个必然会出现的非连续访存发生在指令延迟较低的Cache中，就可以大大降低LDST单元的闲置率，这个时候L1/Shared走进了我们的视野中。
3.2 共享内存
共享内存(Shared Memory)是CUDA编程中的一种软件层概念，对应到实际设备中，就是可编程的L1 Cache。这种片上内存(On-chip Memory)的访存指令延迟非常之低(20cycle左右)，相对于DRAM的内存事务指令来说基本可以忽略不计。那么我们解决此问题的思路就变成了，先将数据整体搬运到共享内存中，在写回时，让非连续访存发生在加载共享内存这个过程中，与此同时仍然保证写回全局内存是连续的，这样带宽利用率会有显著提升。当然反过来也是可以的，就是让非连续访存发生在写入共享内存过程。我们这里以前者为例，尝试优化初版实现。
在实现前，我们先引入双向搬运工(GlobalToShared，SharedToGlobal)以及同步器(Synchronize)，简洁代码的同时减少不必要的运行时计算开销。
（点击查看大图）
（点击查看大图）
（点击查看大图）
接下来进行分块策略的分析，共享内存是初代解决同线程块中不同线程间通信问题的产物。那么基于此，我们把一个线程块的数据全部平移到共享内存中，写回时转置。由于CUDA是以线程束的形式执行指令，因此为了全局内存的访存连续，我们选用以下分块策略：
（点击查看大图）
（点击查看大图）
由于使用了共享内存，所以我们在模板参数中多加入了专门针对于共享内存的存取器 SharedReader 和SharedWriter。此时非连续访存发生在SharedReader中：
（点击查看大图）
我们将线程根据线程束分组，SharedWriter保证每个线程束写共享内存的一行，而SharedReader保证每个线程束读取共享内存的一列。GlobalToSharedWorker和SharedToGlobalWorker两个打工人保证单条数据搬运操作，我们来看一下执行时间：
肉眼可见的速度提升已经出现了，但是刚刚达到cuDNN的速度，还有很多的优化空间。
3.3 Bank Conflict
既然使用了共享内存，那就不得不提到Bank的概念。对于Bank具体感兴趣的同学可以在 Memory Banks 这里了解更多。为了保证内存高带宽，Nvidia在设计L1时，将其分成了32个等大小的内存模块(Memory Modules)，并将其称为(Banks)。共享内存访存有以下3个机制：
1.当一个线程束中的每个线程都访问不同bank时，会触发并行化访存(Parallel Access)机制
2.当一个线程束中的多个线程访问相同bank内的不同地址时，会触发串行访存(Serial Access)机制
3.当一个线程束中的多个线程访问相同bank内的相同地址时，会触发广播访存(Broadcast Access)机制我们注意到了当第二种情况发生时，线程束的32个线程会从完全并发变成部分排队阻塞。我们先来看看3.2中的实现是否有bank conflict。先来看看数据在共享内存中的排布方式(Pattern)：
（点击查看大图）
我们能看到写入共享内存时，确实不存在bank conflict。但是在读取共享内存时，每个warp的32个线程全部读取8个bank中不同的4个地址，因此在理论上会造成最影响性能的8 way bank conflict。用Nsys来验证一下：
正如我们所论述的一样，在共享内存的加载环节出现了8 way bank conflict，冲突数刚好是加载指令数的7倍。那么如何解决呢，其实具体解决方案千变万化，但万变不离其宗的是：通过在特定位置padding的方式，让一个线程束同时访问的共享内存地址所在bank错位开，一般情况下每行padding的总bank数是一个与32互质的数，比如下述方案就是一种解决办法：
（点击查看大图）
通过padding，我们刚巧让同一个线程束访问共享内存的32个bank。那么我们借助这种padding方式实现一下Kernel：修改一下SharedMemShape的列大小，每行padding一个bank就是4字节。
（点击查看大图）除了改变共享内存的列数，其他没有变化。我们通过NCU来验证一下是否解决了bank conflict问题：
再来看一下解决bank conflict后的执行速度：
相对于初版，我们已经获得了4倍的性能提升，而且此时也已经超越了CUTLASS的通用转置算子，但是我们的征程还远未结束。
04向量化访存与PTX优化
其实常规的Kernel优化到上一步解决完bank conflict就已经是工业场景可用的高性能算子了，接下来的优化点主要是SIMT嵌套向量化操作。在此之前需要先介绍一些背景知识。
4.1 向量化访存指令集简介
对于N卡来说，L1/Shared的CacheLine是128Bytes，显然上述优化每个warp仅仅搬运了32Bytes，还远没有达到设备的理论上限，因此理论上让单个线程搬运连续的4Bytes数据，也就是一次搬运4个int8，在理论上可以占满L1/Shared的一次内存事务，这个时候我们需要用到向量化访存。在Turing架构之前的N卡，从全局内存搬运数据到共享内存需要寄存器的介入，也就是说数据真实的传输流程是：
写成代码就是这种形式，对于C/C++内嵌汇编语法不熟悉的同学可以参考 Inline PTX Assembly 学习，不同编译器间大同小异：
（点击查看大图）
需要额外注意的是，共享内存的地址并不是全局同步地址(Generic Address)，因此在使用共享内存地址读取或写入数据前，要经过一次内置函数__cvta_generic_to_shared，当然也可以自己手写PTX：
（点击查看大图）
另外，向量化搬运的单次最大容量为128bit，也就是.v4.u32或.v2.u64，超过上限时汇编器会报错。register关键字仅是提高代码易读性，现代编译器O2选项开启后已不会再理会此种建议性关键字。不难发现，数据在L1 Cache和寄存器之间打圈圈。Nvidia也注意到了这个问题，因此从Ampere架构开始，从全局内存到共享内存，PTX提供了新的指令，减少了打圈圈的过程。也就是说，
代码变成了如下：
（点击查看大图）
这里需要注意的是，当单条指令搬运字节数非16时，只能用.ca qualifier，满16时可以用.cg qualifier。具体差异及其他qualifier的作用可以参考 PTX文档 。新指令cp.async有两种实现机制：分组异步(Async-group mechanism)和基于内存屏障(Mbarrier-based mechanism)，不同的机制使用方法不同，由于第二种机制过于复杂，我们这里仅介绍第一种机制的使用方法。__syncthreads()函数并不帮忙设置组屏障或内存屏障，因此我们需要自己控制屏障的粒度。还是同理做简单封装：
（点击查看大图）
4.2 全局内存到共享内存的单程向量化
我们先来看看将全局内存中的数据搬运到共享内存时使用向量化指令的速度如何。我们不需要改变3.3中的分块规则，只是将ElmentsPerAccess提升到4，当然附之而来的是每个线程块和共享内存的大小增加。也就是单次单线程从全局内存搬运4字节到共享内存，其他模板不变，Kernel实现如下：
（点击查看大图）
先来看看速度如何：
4.3 全局内存到共享内存的双向向量化
对于现阶段的设备来说，除开Hopper架构，共享内存写回到全局内存必然经过一下路径：
看到这个过程，我们自然而然想到了，一次性加载4次u8到4个8bit寄存器(并非真实硬件)中，再向量化写回。先来改造一下SharedToGlobalWorker。
（点击查看大图）
我们看一下共享内存写回全局内存时的真实排布：
（点击查看大图）
所以理论上32线程分128条指令访问分布于32个bank中的128个数，必然会出现4-way bank conflict。先看下Kernel实现：
（点击查看大图）
由于一个线程束要加载32行4列数据，因此我们仅修改了写回过程中的偏移量计算。接下来我们检查Nsys，看一下bank conflict是否是3倍数量：
我们惊讶的发现Nsys的统计和我们的理论并不一致，竟然没有bank conflict。难道是哪一个环节出了问题？没错，确实是现代编译器在内部搞的鬼。我们使用如下指令来看一下生成的fatbin文件内容，或者叫SASS码文件，就是可执行文件的内容：
打开SASS码文件，这里我只截取了写回过程的最开始部分，已经足够发现问题，也就是BAR.SYNC 0x0之后的SASS码：
（点击查看大图）
在这段SASS码中，前面如我们所愿以u8形式加载共享内存的数据到寄存器中，（这里多提一句，由于设备端不具备CPU式的指令乱序执行能力，因此设备执行效率非常依赖N自家的ptxas和fatbinary给出的优化后的指令执行顺序，我们这里看到的加载了12次才执行prmt指令也是拜其所赐），但是第0550行发生了一个我们预期之外的事情，R13和R22被取出了有效位数并合并了到了R2寄存器中，第05b0行也出现了此情况。其实这就是现代编译器的优化，已经极大地提高了代码效率的下限，不用着急，我们在下一节中会深入讨论这种优化，编译器帮我们做了什么。先来看看这时Kernel的执行速度如何：
性能提升已经甚微，此时的带宽利用率已经达到85%，说明基本达到了向量化优化的瓶颈了。我们接下来看看编译器做了什么优化。
05寄存器转置
在第4小节中，写回过程中的向量化搬运会导致有必要深究一下加载共享内存到写回全局内存这一阶段，编译器帮我们做了一些事情消除了理应存在4 way bank conflict。我们有必要深究一下到底是何种优化，既保证不存在bank conflict又可以向量化写回全局内存。
回忆3.3中讨论过的共享内存的3种访存机制，既然bank conflict导致的串行访存无法优化成并行访存，那么我们干脆放弃这条路，尝试是否有广播机制来避免掉串行访存。如果我们单个线程正常加载满4个bank的32bit数据，然后从中提取出我们想要的那4个8bit数据，再重新组装成一个新的32bit数据，这个问题不就被完美解决了么？借助于prmt这个字节拆解打包指令，寄存器转置由此成为可能，这也就是编译器帮我们做的优化，通过广播机制避免bank conflict。
（点击查看大图）
当i，i + 8, i + 16, i + 24 (0 <= i < 8)号线程访问第 4 i ~ 4 (i + 1) 行时，我们把全部 4 4 大小的s8数据以4个u32的形式全部读出，使用prmt指令组装成对应所需的第0，1， 2， 3列数据，打包成一个u32写回全局内存，就刚好实现了向量化写回的过程。之所以叫它寄存器转置，是因为本质上，在4 4的小矩阵中转置后再写回对应位置。让我们先来了解一下prmt(Permute)指令的工作原理和大小端相关的背景知识。
prmt指令是在两个32bit寄存器的总共8个字节中，取出任意的2个字节并打包到一个32bit寄存器的指令。当我们从共享内存中加载4个32bit，并从中取出固定列数的4个字节打包成一个32bit的过程可以描述为下图：
（点击查看大图）
由于在共享内存中，数据是以小端形式存储的，也就是高地址对应高位，低地址对应地位。那么在前两次合并时，对于取0，1，2，3列，需要的ctl数分别为0x0040，0x0051，0x0062和0x0073。第三次合并就是固定的0x0145。当然也有其他很多种拆分方式，我们这里就以此种拆分为例，写成代码就是如下形式：
（点击查看大图）
需要注意，64bit模式需要先将设备的Bank Size开成8Bytes大小。基于寄存器转置的Kernel实现如下：
（点击查看大图）
我们现在再来看看SASS码，同样只截取写回过程：
（点击查看大图）
有没有发现，当我们帮助编译器做了基于广播机制的寄存器转置优化后，编译器再一次智能地帮我们的共享内存访存指令进行了合并。每条指令已经进化到了64bit或128bit，这得益于我们最开始的warp分块策略。我们再来看看目前的性能：
至此，带宽方向的优化已基本拉满，达到了4090的87%左右，剩下的性能瓶颈我们就需要打开Nsys针对性分析了。对于我们这种无计算型算子，观察Nsys给出的感叹号，仅剩Warp Stall Statistics一项，点开它：
（点击查看大图）
我们注意到Stall Long Scoreboard是目前的性能瓶颈，换回大白话讲就是，写回过程需要等待加载过程全部结束后方可进行，这样就导致BAR.SYNC时，有非常多的线程被迫驻留等待全部128线程执行完毕，我们知道设备端的写入和写出是可以同时进行而不会互相阻塞的，也就是说一条传输线可以同时进行双向传输，那么理论上还有最后一种优化 ------ pipeline buffer 可以有效提升性能瓶颈，让写入和写回彻底异步化。由于本文篇幅限制，pipeline buffer就交给大家去探索啦！
06总结
6.1 与cuDNN，CUTLASS的性能对比
（点击查看大图）
6.2 CUDA访存优化策略
我们通过转置算子的优化，一步一步的揭开了CUDA优化的面纱。算子一般情况下分为访存和计算两个部分，由于现代设备计算指令的延迟远远低于访存类指令的延迟，因此访存是我们优化的重点，针对于计算量较大的算子，我们可以将计算部分拆分到访存指令之间，已达到延迟覆盖(Latency Hidding)的目的。CUDA访存优化一般分为以下几个步骤：
1.任务拆分，将较大任务根据不同的层级进行分块，用并行编程实现初版功能
2.检查访存是否连续，尽量保证在读写全局内存时，一个线程束访问一块连续的内存地址
3.将无法连续访存的部分，或是使用频繁且数据量较大的部分搬运至延迟更低的片上内存中(共享内存)
4.检查共享内存的访存是否存在bank conflict问题
5.向量化访存，以达到SIMT嵌套向量化操作
6.除了共享内存，还有访存延迟更低的寄存器，同一个线程块内线程间的通讯还可以用shfl.sync等指令
再配以模板和宏的灵活使用，将一些不必要的计算从运行时搬至编译期甚至是预处理期，并简化代码结构。这样一个较高性能的算子就此诞生。