Shared memory！CUDA数据拷贝速度拉满~

作者 | Swaghe  编辑 | 自动驾驶之心

原文链接：https://zhuanlan.zhihu.com/p/694779147

点击下方卡片，关注“自动驾驶之心”公众号

戳我-> 领取自动驾驶近15个方向学习路线

>>点击进入→自动驾驶之心『模型部署』技术交流群

本文只做学术分享，如有侵权，联系删文

最近遇到了一个问题，拷贝数据的时候经过shared memory的带宽会快于没有经过shared memory，使用shared memory的具体代码如下， 拷贝的数据大小是4096 * 4096个double数据（128MiB），blocksize是（1024，1，1），gridsize是（4096*4096/1024/2，1，1）。

// blocksize=(1024，1，1), gridsize=(4096*4096/1024/2，1，1)
__global__ void copySheme(double * MatA,double* MatB)
{
  __shared__ double tile[1024*2];
  size_t idx=threadIdx.x+blockDim.x * blockIdx.x * 2;

  tile[threadIdx.x]=MatA[idx];
  tile[threadIdx.x+blockDim.x]=MatA[idx+blockDim.x];

  MatB[idx]=tile[threadIdx.x];
  MatB[idx+blockDim.x]=tile[threadIdx.x+blockDim.x];
}

没有使用shared memory的代码如下，此时，blocksize和gridsize不变。

// blocksize=(1024，1，1), gridsize=(4096*4096/1024/2，1，1)
__global__ void copyRow(double * in,double * out)
{
  size_t idx=threadIdx.x+blockDim.x * blockIdx.x * 2;

  out[idx]  = in[idx];
  out[idx+blockDim.x]  = in[idx+blockDim.x];

}

在A800 80GB上，用nsys 测得copySheme的平均运行时间为153495.7ns。用nsys 测得copyRow的平均运行时间为157423.7ns。

最开始不能理解，因为认为数据经过shared memory多出一层开销，为什么速度反而会变快呢，后来参考英伟达官方论坛的解释：

forums.developer.nvidia.com/t/why-copy-using-shared-memory-is-faster-than-direct-copy/70257/3

通过shared memory，使得对DRAM的load和store分开了，不是交错进行load和store，这使得硬件效率更高。通过查看对应kernel的SASS，也确实发现copySheme是先执行两条load指令取数，再执行两条store指令存数。而copyRow则是load和store交替执行。

但是在阅读 copySheme的SASS代码也发现，从DRAM取数据放入shared memory 再存储到DRAM的数据路径大致是：DRAM->Reg->Shared memory->DRAM。因为没有这么一条指令可以直接将数据从DRAM取出并放入Shared memory。类似的，copyRow的数据路径大致是：DRAM->Reg->DRAM。这也说明了copySheme确实多一层Shared memory的开销。

因此，能不能综合以上两种方法的优点呢，即让数据的拷贝没有shared memory的开销，同时把多条load和store分成两部分执行，即先执行完所有的load指令，再执行所有的store指令。因此我想到了下面这种方法：

// blocksize=(1024，1，1), gridsize=(4096*4096/1024/2，1，1)
__global__ void copyReg(double * in,double * out)
{ 
    int idx=threadIdx.x+blockDim.x*blockIdx.x*2;
    double tmp0 = in[idx];
    double tmp1 = in[idx+blockDim.x];
    
    out[idx] = tmp0 ;
    out[idx+blockDim.x] = tmp1;
}

通过显式的声明临时变量来代替shared memory的功能，因为变量的值是存储到寄存器中的。很不幸，测出来的时间是153583.1ns ，略高于copySheme的 153495.7ns，而且这个时间很稳定，总是略高于，我对比了两者的SASS，copyReg确实达到了我想要的执行方式，指令条数也小于copySheme，但是时间仍然高于 copySheme。这是让我疑惑的地方 。

copyReg的SASS如下：

copyReg(double*, double*):
 MOV R1, c[0x0][0x28] 
 S2R R0, SR_CTAID.X 
 HFMA2.MMA R9, -RZ, RZ, 0, 4.76837158203125e-07 
 ULDC.64 UR4, c[0x0][0x118] 
 S2R R3, SR_TID.X 
 SHF.L.U32 R0, R0, 0x1, RZ 
 IMAD R0, R0, c[0x0][0x0], R3 
 IMAD.WIDE R2, R0.reuse, R9, c[0x0][0x160] 
 IADD3 R8, R0, c[0x0][0x0], RZ 
 IMAD.WIDE.U32 R4, R8, R9.reuse, c[0x0][0x160] 
 LDG.E.64 R2, [R2.64] 
 LDG.E.64 R4, [R4.64] 
 IMAD.WIDE R6, R0, R9, c[0x0][0x168] 
 IMAD.WIDE.U32 R8, R8, R9, c[0x0][0x168] 
 STG.E.64 [R6.64], R2 
 STG.E.64 [R8.64], R4 
 EXIT

copySheme的SASS如下：

copySheme(double*, double*):
 MOV R1, c[0x0][0x28] 
 S2R R0, SR_CTAID.X 
 HFMA2.MMA R11, -RZ, RZ, 0, 4.76837158203125e-07 
 ULDC.64 UR4, c[0x0][0x118] 
 S2R R15, SR_TID.X 
 SHF.L.U32 R0, R0, 0x1, RZ 
 IMAD R0, R0, c[0x0][0x0], R15 
 IMAD.WIDE.U32 R2, R0, R11, c[0x0][0x160] 
 IMAD.WIDE.U32 R4, R11, c[0x0][0x0], R2 
 LDG.E.64 R2, [R2.64] 
 LDG.E.64 R4, [R4.64] 
 SHF.L.U32 R6, R15, 0x3, RZ 
 IMAD.WIDE.U32 R8, R0, R11, c[0x0][0x168] 
 IMAD R13, R11.reuse, c[0x0][0x0], R6 
 IMAD.WIDE.U32 R10, R11, c[0x0][0x0], R8 
 STS.64 [R15.X8], R2 
 STS.64 [R13], R4 
 LDS.64 R6, [R15.X8] 
 STG.E.64 [R8.64], R6 
 STG.E.64 [R10.64], R4 
 EXIT

如果确实是按猜想所说的将load和store分批次执行对硬件更友好，能提升访存性能，如果我提升每批次执行的load和store的数量，会不会使得性能更好呢，我们进一步优化copyReg如下：

// blocksize=(1024，1，1), gridsize=(4096*4096/1024/unroll_size ，1，1)
__global__ void copyReg_v2(double * in,double * out)
{ 
    size_t unroll_size = 16;
    int cidx=threadIdx.x+blockDim.x*blockIdx.x*unroll_size;
    int idx = cidx;
    // 这里实际执行的时候要改成具体的数据，这里只是方便理解
    double tmp[unroll_size];
    
    for (int i = 0; i < unroll_size; i++){
      tmp[i] = in[idx + i * blockDim.x];
    }

    for (int i = 0; i < unroll_size; i++){
      out[idx + i * blockDim.x] = tmp[i];
    }
}

通过控制变量unroll_size的大小来控制每批次执行的load和store的次数，即先执行 unroll_size次load，再执行unroll_size次store，很遗憾，这样并没有逐步提升访存性能。大致实验结果如下：

unroll_size = 16， 时间 = 156340.2ns；

unroll_size = 8， 时间 = 154622.2ns；

unroll_size = 4， 时间 = 154866.8ns；

unroll_size = 2， 时间 = 153569.1；

关于这个现象的一个猜测就是提升unroll_size的时候，寄存器的使用量也大大的提升了，目前一个block最多的寄存器使用量是65536，unroll_size=16的时候，一个线程要使用40个寄存器，一个block使用40960个寄存器。这个信息可以通过加上编译选项-Xptxas=-v看到。

目前现象大致为文章描述的情况，大家有什么见解也可以在评论区一起讨论。

投稿作者为『自动驾驶之心知识星球』特邀嘉宾，欢迎加入交流！

① 全网独家视频课程

BEV感知、BEV模型部署、BEV目标跟踪、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、车道线检测、轨迹预测、在线高精地图、世界模型、点云3D目标检测、目标跟踪、Occupancy、cuda与TensorRT模型部署、大模型与自动驾驶、Nerf、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频（扫码即可学习）

网页端官网：www.zdjszx.com

② 国内首个自动驾驶学习社区

国内最大最专业，近3000人的交流社区，已得到大多数自动驾驶公司的认可！涉及30+自动驾驶技术栈学习路线，从0到一带你入门自动驾驶感知（2D/3D检测、语义分割、车道线、BEV感知、Occupancy、多传感器融合、多传感器标定、目标跟踪）、自动驾驶定位建图（SLAM、高精地图、局部在线地图）、自动驾驶规划控制/轨迹预测等领域技术方案、大模型、端到端等，更有行业动态和岗位发布！欢迎扫描下方二维码，加入自动驾驶之心知识星球，这是一个真正有干货的地方，与领域大佬交流入门、学习、工作、跳槽上的各类难题，日常分享论文+代码+视频

③【自动驾驶之心】技术交流群

自动驾驶之心是首个自动驾驶开发者社区，聚焦感知、定位、融合、规控、标定、端到端、仿真、产品经理、自动驾驶开发、自动标注与数据闭环多个方向，目前近60+技术交流群，欢迎加入！

自动驾驶感知：目标检测、语义分割、BEV感知、毫米波雷达视觉融合、激光视觉融合、车道线检测、目标跟踪、Occupancy、深度估计、transformer、大模型、在线地图、点云处理、模型部署、CUDA加速等技术交流群；

多传感器标定：相机在线/离线标定、Lidar-Camera标定、Camera-Radar标定、Camera-IMU标定、多传感器时空同步等技术交流群；

多传感器融合：多传感器后融合技术交流群；

规划控制与预测：规划控制、轨迹预测、避障等技术交流群；

定位建图：视觉SLAM、激光SLAM、多传感器融合SLAM等技术交流群；

三维视觉：三维重建、NeRF、3D Gaussian Splatting技术交流群；

自动驾驶仿真：Carla仿真、Autoware仿真等技术交流群；

自动驾驶开发：自动驾驶开发、ROS等技术交流群；

其它方向：自动标注与数据闭环、产品经理、硬件选型、求职面试、自动驾驶测试等技术交流群；

扫码添加汽车人助理微信邀请入群，备注：学校/公司+方向+昵称（快速入群方式）

④【自动驾驶之心】全平台矩阵