LeetCUDA-学习记录-点积(dot product)

LeetCUDA-学习记录

---

---

前言

本系列文章主要是记录作者本人学习LeetCUDA项目的笔记以及作者本人对该项目中对于一系列算子设计的思路进行理解分析与优化分析，有所错误望大佬们指正~

---

一、矩阵转置(mat-transpose)

二、规约计算(reduce)

后面持续更新中~~OVO

三、点积(dot product)

点积(dot product) 顾名思义就是向量的对应元素乘积再求和。个人感觉点积的设计，在对于矩阵乘法方面有着优化指导（比如一行一列向量乘积求和）。
本文主要是对LeetCUDA项目关于dot_product点积的思路设计进行从为什么的角度上理解分析，且由于设计算子针对的数据类型不同但是思路相同，因此本文只分析其中的32数据类型向量的点积思路。

项目设计思路

按照LeetCUDA项目，个人认为主要是从float4向量化+shuffle寄存器指令计算两个角度进行优化。代码如下：

#define FLOAT4(value) (reinterpret_cast<float4 *>(&(value))[0])

template <const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_sum_f32(float val){
    #pragma unroll
    for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1){
        val += __shfl_xor_sync(0xffffffff, val, mask);
    }
    return val;
}

template <const int NUM_THREADS_PER_BLOCK = 256 / 4>
__global__ void dot_prod_f32x4_f32_kernel(float *a, float *b, float *y, int N){
    int tid = threadIdx.x;
    int idx = (blockIdx.x * NUM_THREADS_PER_BLOCK + tid) * 4;
    // 一个block中总的warp数量
    constexpr int NUM_WARPS_PER_BLOCK = (NUM_THREADS_PER_BLOCK + WARP_SIZE - 1) / WARP_SIZE;
    // 共享内存存储每个warp计算后的结果
    __shared__ float reduce_smem[NUM_WARPS_PER_BLOCK];

    float4 reg_a = FLOAT4(a[idx]);
    float4 reg_b = FLOAT4(b[idx]);
    float prod = (idx < N) ? (reg_a.x * reg_b.x + reg_a.y * reg_b.y + reg_a.z * reg_b.z + reg_a.w * reg_b.w) : 0.f;

    int warp_idx = tid / WARP_SIZE;
    int warp_lane = tid % WARP_SIZE;
	// 一个warp中计算规约求和
    prod  = warp_reduce_sum_f32<WARP_SIZE>(prod);
    if(warp_lane == 0) reduce_smem[warp_idx] = prod;
    __syncthreads(); //这个同步是不同warp之间的同步，而同一warp中的线程是隐式同步的
    //上面显示同步的使用，就是将所有warp计算的结果都保存到了每个warp中第一线程管理的寄存器中
    // 并且都保存到共享内存当中，以达到不同warp之间的结果读取
	
	// 这里将共享内存上的计算结果读取到一个warp中, 以WARP数量为使用的线程数，来分别读取数据到同一WARP的寄存器中
    prod = (warp_lane < NUM_WARPS_PER_BLOCK) ? reduce_smem[warp_lane] : 0.f;
    // 由于不同warp计算的结果都保存到同一warp的寄存器中，因此继续采用寄存器之间的规约求和，得到所有warp计算结果之和
    if(warp_idx == 0) prod = warp_reduce_sum_f32<NUM_WARPS_PER_BLOCK>(prod);
    
	// 为什么使用tid==0，是因为上述warp之间结果求和到存储到tid==0的寄存器中，当tid==0时，prod读取出来的值就是总和
    if(tid == 0) atomicAdd(y, prod); // 将每个block计算的结果统一使用原子加法操作加入到同一个地址当中


}

首先对核函数dot_prod_f32x4_f32_kernel中逐行分析以及为什么这样写。

#define FLOAT4(value) (reinterpret_cast<float4 *>(&(value))[0])

int idx = (blockIdx.x * NUM_THREADS_PER_BLOCK + tid) * 4;
constexpr int NUM_WARPS_PER_BLOCK = (NUM_THREADS_PER_BLOCK + WARP_SIZE - 1) / WARP_SIZE;
__shared__ float reduce_smem[NUM_WARPS_PER_BLOCK];
float4 reg_a = FLOAT4(a[idx]);
float4 reg_b = FLOAT4(b[idx]);
float prod = (idx < N) ? (reg_a.x * reg_b.x + reg_a.y * reg_b.y + reg_a.z * reg_b.z + reg_a.w * reg_b.w) : 0.f;

从float4向量化角度思考：期望每个线程能处理4个float32数据，因此需要的总线程数则是 N / 4 (其中N是向量的大小)。idx访问指定地址也是以4为stride进行跳跃访问。
从shuffle寄存器指令计算角度思考：一般来说，在同一warp当中，线程计算结果后写入寄存器中，然而当另一个线程想要获取该线程计算的结果时，必须将计算结果写入到共享内存当中，然后再进行读取以获取计算结果，这个过程需要四次读写操作，分别从寄存器中读，到写入共享内存，再从共享内存中读，最后写入寄存器。为了减少读写操作也就是访问延迟，nv官方提出了shuffle寄存器指令操作，能够让同一warp中的线程可以互相访问线程写入寄存器位置的值。因此作者创建一个reduce_smem[NUM_WARPS_PER_BLOCK]共享内存，来存储每个warp内部线程计算乘积后规约求和后的结果。也就有了如下代码：

template <const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_sum_f32(float val){
    #pragma unroll
    for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1){
        val += __shfl_xor_sync(0xffffffff, val, mask);
    }
    return val;
}
int warp_idx = tid / WARP_SIZE;
int warp_lane = tid % WARP_SIZE;

prod  = warp_reduce_sum_f32<WARP_SIZE>(prod);
if(warp_lane == 0) reduce_smem[warp_idx] = prod;
__syncthreads(); 

自定义设备函数warp_reduce_sum_f32当中，采用的__shfl_xor_sync函数就是通过mask与线程异或的形式来读取另一个线程的值，思路如下：

因此当warp_lane == 0 时，由于warp内部隐式的同步，当你执行完prod = warp_reduce_sum_f32<WARP_SIZE>(prod);，同一个warp内部线程就已经完成了上述的shuffle操作，并且32个线程计算的结果值求和都保存到了线程0所管理的寄存器地址上的值。从而只要warp_lane == 0 时将计算结果存放到共享内存上就可以了。
这个__syncthreads()则是避免不同warp之间并不同步，从而等待所有warp将其计算结果保存到共享内存当中后，再进行如下操作：

// 这里将共享内存上的计算结果读取到一个warp中, 以WARP数量为使用的线程数，来分别读取数据到同一WARP的寄存器中
// 由于按照一个block最多线程量为1024（不同GPU不同大小），所有warp数量都不会超过32，所以共享内存读取就不会出现bank conflict问题
prod = (warp_lane < NUM_WARPS_PER_BLOCK) ? reduce_smem[warp_lane] : 0.f;

// 由于不同warp计算的结果都保存到同一warp的寄存器中，因此继续采用寄存器之间的规约求和，得到所有warp计算结果之和
if(warp_idx == 0) prod = warp_reduce_sum_f32<NUM_WARPS_PER_BLOCK>(prod);

// 为什么使用tid==0，是因为上述warp之间结果求和到存储到tid==0对应管理的寄存器中，当tid==0时，prod读取出来的值就是总和
if(tid == 0) atomicAdd(y, prod); // 将每个block计算的结果统一使用原子加法操作加入到同一个地址当中

正如注释所说，最后同样减少对共享内存的访存次数，利用寄存器的访问时间比共享内存的短的优势，将共享内存上的计算结果都写入同一个warp当中，让不同线程对应的寄存器地址上的值进行规约求和，从而得到总和的值，再使用atomicAdd()写入。

相关参考资料

atomicAdd()参考资料
shuffle函数参考资料

四、矩阵乘法(gemm)

后面持续更新中~~OVO