转置是深度模型中利用最宽泛的算子之一。各种深度学习框架崛起的明天,它被“隔离”到了框架的底层,不再显式的沉闷于用户背后。谈及转置,就不得不先聊聊卷积的优化算法以及它们所须要的图像数据格式。
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等指令

再配以模板和宏的灵便应用,将一些不必要的计算从运行时搬至编译期甚至是预处理期,并简化代码构造。这样一个较高性能的算子就此诞生。