共计 8405 个字符，预计需要花费 22 分钟才能阅读完成。
转置是深度模型中利用最宽泛的算子之一。各种深度学习框架崛起的明天，它被“隔离”到了框架的底层，不再显式的沉闷于用户背后。谈及转置，就不得不先聊聊卷积的优化算法以及它们所须要的图像数据格式。
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 等指令
再配以模板和宏的灵便应用，将一些不必要的计算从运行时搬至编译期甚至是预处理期，并简化代码构造。这样一个较高性能的算子就此诞生。