实操教程|CUDA WarpReduce 学习笔记

极市平台

共 1501字,需浏览 4分钟

 ·

2022-04-12 02:51

↑ 点击蓝字 关注极市平台

作者丨zzk
来源丨GiantPandaCV
编辑丨极市平台

极市导读

 

CUDA 编程进阶分享,希望这篇博客能帮助读者理解到一些 warp 的使用。 >>加入极市CV技术交流群,走在计算机视觉的最前沿

前言

之前看我司的 如何实现一个高效的Softmax CUDA kernel?多少还是有些细节没有理解,恰好最近要做一个类似的 Reduce+Scale Kernel,原理机制还是比较相似的,所以翻出来重新理解一下。

背景

我们定义这么一个ReduceScale操作:假设Tensor是(N, C),首先在C这个维度计算出 absMax 值,我们记作scale,然后将每一行除以各自 行的scale,并最终输出。一段朴素的numpy代码是这样:

import numpy as np  
  
  
N = 1000  
C = 128  
x = np.random.randn(N, C)  
scale = np.expand_dims(np.max(np.abs(x), axis=1), 1)  
out = x / scale  
print(out.shape)  

BaseLine

这里我们BaseLine是直接调用cub库中的 BlockReduce,一个 threadBlock 处理一行数据,计算出AbsMaxVal,然后再缩放,代码如下:

#include "cuda.h"  
#include "cub/cub.cuh"  
  
constexpr int kReduceBlockSize = 128;  
  
template  
__device__ T abs_func(const T& a) {  
  return abs(a);  
}  
  
  
template  
__device__ T max_func(const T a, const T b) {  
  return a > b ? a : b;  
}  
  
template  
struct AbsMaxOp {  
  __device__ __forceinline__ T operator()(const T& a, const T& b) const {  
    return max_func(abs_func(a), abs_func(b));  
  }  
};  
  
template  
__inline__ __device__ T BlockAllReduceAbsMax(T val) {  
  typedef cub::BlockReduce BlockReduce;  
  __shared__ typename BlockReduce::TempStorage temp_storage;  
  __shared__ T final_result;  
  T result = BlockReduce(temp_storage).Reduce(val, AbsMaxOp());  
  if (threadIdx.x == 0) { final_result = result; }  
  __syncthreads();  
  return final_result;  
}  
  
template  
__global__ void ReduceScaleBlockKernel(T* x, IDX row_size, IDX col_size) {  
  for(int32_t row = blockIdx.x, step=gridDim.x; row < row_size; row+= step){  
    T thread_scale_factor = 0.0;   
    for(int32_t col=threadIdx.x; col < col_size; col+= blockDim.x){  
      IDX idx = row * col_size + col;   
      T x_val = x[idx];  
      thread_scale_factor = max_func(thread_scale_factor, abs_func(x_val));   
    }  
    T row_scale_factor = BlockAllReduceAbsMax(thread_scale_factor);   
    for(int32_t col=threadIdx.x; col < col_size; col+=blockDim.x){  
      IDX idx = row * col_size + col;   
      x[idx] /= row_scale_factor;  
    }  
  }  
}  

参数中 x 是输入数据,row_size是行的数量,col_size是列的大小测试机器是在 A100 40GB,为了让结果区别比较明显,我们将行数设置的比较大,输入形状为(55296*8, 128),启动的线程块数目根据 如何设置CUDA Kernel中的grid_size和block_size?这篇文章来指定,这里比较粗暴的设置为(55296, 128),数据类型为 Float,然后我们看下ncu的结果:

主要有这几个指标,耗时为577.95us,吞吐量为 748.78Gb/s下面我们就根据 Softmax 优化那篇文章所提及的点来逐步分析:

优化1 数据Pack

在之前的 高效、易用、可拓展我全都要:OneFlow CUDA Elementwise 模板库的设计优化思路 里很详细的描述了如何做向量化读写,cuda里最大支持 128bit的读写,那么在数据类型为 Float 时,我们即可以将连续的4个 Float 打包到一起,一次性读写,提升吞吐。有了解过这方面的读者应该就反应过来,诶 CUDA 里 不是刚好有一个类型叫 float4 就是干这件事的么,没错,但是为了更灵活的支持其他数据类型的向量化,我们利用union共享空间的特性实现了一个 Pack 类:

template  
struct GetPackType {  
  using type = typename std::aligned_storage::type;  
};  
  
template  
using PackType = typename GetPackType::type;  
  
template  
union Pack {  
  static_assert(sizeof(PackType) == sizeof(T) * N, "");  
  __device__ Pack() {  
    // do nothing  
  }  
  PackType storage;  
  T elem[N];  
};  

优化2 数据缓存

整个算子逻辑是需要读取一遍数据,计算scale,然后再读取一遍数据,用scale进行缩放。很显然这里我们读取了两遍数据,而数据是放在 Global Memory,带宽比较低,会带来读取耗时。

一个很自然的想法是缓存到寄存器/Shared Memory中。由于这里我们只实现 WarpReduce 版本,所以我们是缓存到寄存器(其他版本可以参考开头的优化 Softmax 文章)中,减少一次对 Global Memory 的读取。

template  
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {  
    // ...  
    T buf[cols_per_thread];  
    // ...  

优化3 使用Warp处理一行数据

相较 BaseLine,我们这里使用 warp 作为 Reduce 的单位进行操作,首先我们简单看下 WarpReduce 的实现。

template  
struct AbsMaxOp {  
  __device__ __forceinline__ T operator()(const T& a, const T& b) const {  
    return max_func(abs_func(a), abs_func(b));  
  }  
};  
  
template  
__inline__ __device__ T WarpAbsMaxAllReduce(T val){  
    for(int lane_mask = kWarpSize/2; lane_mask > 0; lane_mask /= 2){  
        val = AbsMaxOp()(val, __shfl_xor_sync(0xffffffff, val, lane_mask));   
    }  
    return val;   
}  

这段代码在别的 BlockReduce 也经常看到,他是借助 __shfl_xor_sync 来实现比较,shuffle 指令允许同一线程束的两个线程直接读取对方的寄存器。

T __shfl_xor_sync(unsigned mask, T var, int laneMask, int width=warpSize);  

其中 mask 是对线程的一个掩码,我们一般所有线程都要参与计算,所以 mask 是 0xffffffffvar 则是寄存器值,laneMask 则是用来做按位异或的掩码

这里引入一个概念叫 Lane,它表示线程束中的第几号线程

示意图如下:

当 laneMask = 16 时,其二进制为 0001 0000,然后线程束每个线程与 laneMask 做异或操作如:

  • 0000 0000 xor 0001 0000 = 0001 0000 = 16
  • 0000 0001 xor 0001 0000 = 0001 0001 = 17
  • 0000 0010 xor 0001 0000 = 0001 0010 = 18

以此类推,最终得到一个 Warp 中的 absmax 值。接下来我们开始写Kernel,模板参数分别为:

  • T 数据类型
  • IDX 索引类型
  • pack_size pack数,比如float可以pack成4个,那对应pack_size=4
  • cols_per_thread 每个线程需要处理的元素个数,比如一行大小是128,而我们一个warp有32个线程,那么这里就是128/32 = 4
template  
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {  
    // ...      
}  

跟BaseLine一样,我们block大小还是设置为128个线程,一个warp是32个线程,所以我们一个block可以组织成(32, 4),包含4个warp。

根据这个层级划分,我们可以计算出:

  • global_thread_group_id 当前warp的全局index
  • num_total_thread_group warp的总数量
  • lane_id 线程束内的线程id
  • num_packs pack的数目,即每个线程需要处理的元素个数 / pack_size
const int32_t global_thread_group_id = blockIdx.x * blockDim.y + threadIdx.y;   
    const int32_t num_total_thread_group = gridDim.x * blockDim.y;   
    const int32_t lane_id = threadIdx.x;   
    using LoadStoreType = PackType;  
    using LoadStorePack = Pack;  
    T buf[cols_per_thread];   
    constexpr int num_packs = cols_per_thread / pack_size;  

由于存在启动的warp的数量小于行的数量,所以我们要引入一个 for 循环。假设我们 cols = 256,那么线程束里的每个线程需要处理 256 /32 = 8个元素,而4个float可以pack到一起,所以我们线程束里的每个线程要处理2个pack,因此也要引入一个关于 num_packs 的 for 循环,以保证整一行都有被读取到:

一次性读取到一个 pack 后,我们再一个个放到寄存器当中缓存起来,并计算线程上的 AbsMaxVal。

for(IDX row_idx = global_thread_group_id; row_idx < row_size; row_idx += num_total_thread_group){  
        T thread_abs_max_val = 0.0;   
        for(int pack_idx = 0; pack_idx < num_packs; pack_idx++){  
            const int32_t pack_offset = pack_idx * pack_size;   
            const int32_t col_offset = pack_idx * kWarpSize * pack_size + lane_id * pack_size;   
            const int32_t load_offset = (row_idx * col_size + col_offset) / pack_size;   
            LoadStorePack load_pack;   
            load_pack.storage = *(reinterpret_cast(x)+ load_offset);   
            #pragma unroll   
            for(int i = 0; i < pack_size; i++){  
                buf[pack_offset] = load_pack.elem[i];   
                thread_abs_max_val = max_func(thread_abs_max_val, abs_func(buf[pack_offset]));  
            }   
        }  

接着我们调用 WarpAbsMaxAllReduce 进行reduce,获得线程束中的 AbsMaxVal,并对缓存的数据进行数值缩放。

T warp_max_val = WarpAbsMaxAllReduce(thread_abs_max_val);   
        #pragma unroll  
        for (int col = 0; col < cols_per_thread; col++) {  
            buf[col] = buf[col] / warp_max_val;  
        }  

最后跟一开始读取类似,我们将寄存器里的值再写回去,相关索引的计算逻辑都是一致的:

for(int pack_idx = 0; pack_idx < num_packs; pack_idx++){  
            const int32_t pack_offset = pack_idx * pack_size;   
            const int32_t col_offset = pack_idx * pack_size * kWarpSize + lane_id * pack_size;   
            const int32_t store_offset = (row_idx * col_size + col_offset) / pack_size;   
            LoadStorePack store_pack;   
            #pragma unroll   
            for(int i = 0; i < pack_size; i++){  
                store_pack.elem[i] = buf[pack_offset + i];   
            }   
            *(reinterpret_cast(x)+ store_offset) = store_pack.storage;   
        }  

完整代码如下:

template  
__inline__ __device__ T WarpAbsMaxAllReduce(T val){  
    for(int lane_mask = kWarpSize/2; lane_mask > 0; lane_mask /= 2){  
        val = AbsMaxOp()(val, __shfl_xor_sync(0xffffffff, val, lane_mask));   
    }  
    return val;   
}  
  
template  
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {  
    const int32_t global_thread_group_id = blockIdx.x * blockDim.y + threadIdx.y;   
    const int32_t num_total_thread_group = gridDim.x * blockDim.y;   
    const int32_t lane_id = threadIdx.x;   
    using LoadStoreType = PackType;  
    using LoadStorePack = Pack;  
    T buf[cols_per_thread];   
    constexpr int num_packs = cols_per_thread / pack_size;  
    for(IDX row_idx = global_thread_group_id; row_idx < row_size; row_idx += num_total_thread_group){  
        T thread_abs_max_val = 0.0;   
        for(int pack_idx = 0; pack_idx < num_packs; pack_idx++){  
            const int32_t pack_offset = pack_idx * pack_size;   
            const int32_t col_offset = pack_idx * kWarpSize * pack_size + lane_id * pack_size;   
            const int32_t load_offset = (row_idx * col_size + col_offset) / pack_size;   
            LoadStorePack load_pack;   
            load_pack.storage = *(reinterpret_cast(x)+ load_offset);   
            #pragma unroll   
            for(int i = 0; i < pack_size; i++){  
                buf[pack_offset] = load_pack.elem[i];   
                thread_abs_max_val = max_func(thread_abs_max_val, abs_func(buf[pack_offset]));  
            }   
        }  
        T warp_max_val = WarpAbsMaxAllReduce(thread_abs_max_val);   
        #pragma unroll  
        for (int col = 0; col < cols_per_thread; col++) {  
            buf[col] = buf[col] / warp_max_val;  
        }  
        for(int pack_idx = 0; pack_idx < num_packs; pack_idx++){  
            const int32_t pack_offset = pack_idx * pack_size;   
            const int32_t col_offset = pack_idx * pack_size * kWarpSize + lane_id * pack_size;   
            const int32_t store_offset = (row_idx * col_size + col_offset) / pack_size;   
            LoadStorePack store_pack;   
            #pragma unroll   
            for(int i = 0; i < pack_size; i++){  
                store_pack.elem[i] = buf[pack_offset + i];   
            }   
            *(reinterpret_cast(x)+ store_offset) = store_pack.storage;   
        }  
    }  
}  

这里我们方便测试,调用的时候就直接写死一些模板参数

constexpr int cols_per_thread = 128 / kWarpSize;   
ReduceScaleWarpKernel<float, int32_t, 4, cols_per_thread><<<55296, block_dim>>>(device_ptr, row_size, col_size);  

最后我们看一下 ncu 的结果:

吞吐量达到了1.3T,时间位333us,相比 BaseLine 快了 73 %。

总结

还有更多特殊情况可以参考 Softmax 优化的代码,这里仅实现了第一个 Warp 计算方式。我感觉看着还行,真自己写起来理解还是有点困难的,希望这篇博客能帮助读者理解到一些 warp 的使用。


公众号后台回复“画图模板”获取90+深度学习画图模板~

△点击卡片关注极市平台,获取最新CV干货
极市干货
数据集资源汇总:90+深度学习开源数据集整理|包括目标检测、工业缺陷、图像分割等多个方向
实操教程实操教程|Pytorch转ONNX详解一文解决样本不均衡(全)
CVPR 2022:CVPR'22 最新132篇论文分方向整理CVPR'22 最新106篇论文分方向整理


CV技术社群邀请函 #

△长按添加极市小助手
添加极市小助手微信(ID : cvmart4)

备注:姓名-学校/公司-研究方向-城市(如:小极-北大-目标检测-深圳)


即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群


每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~


觉得有用麻烦给个在看啦~  
浏览 19
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报