点击上方“计算机视觉工坊”,选择“星标”
干货第一时间送达
作者丨zzk
来源丨 GiantPandaCV
前言
之前看我司的 如何实现一个高效的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<typename T>
__device__ T abs_func(const T& a) {
return abs(a);
}
template<typename T>
__device__ T max_func(const T a, const T b) {
return a > b ? a : b;
}
template<typename T>
struct AbsMaxOp {
__device__ __forceinline__ T operator()(const T& a, const T& b) const {
return max_func(abs_func(a), abs_func(b));
}
};
template<typename T>
__inline__ __device__ T BlockAllReduceAbsMax(T val) {
typedef cub::BlockReduce<T, kReduceBlockSize> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
__shared__ T final_result;
T result = BlockReduce(temp_storage).Reduce(val, AbsMaxOp<T>());
if (threadIdx.x == 0) { final_result = result; }
__syncthreads();
return final_result;
}
template<typename T, typename IDX>
__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<T>(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的结果:
![c4dbbc2d9aa6f06b4307161a4f153153.png](https://i-blog.csdnimg.cn/blog_migrate/8a04fee841f95ba792239244ffdbd9e7.png)
主要有这几个指标,耗时为577.95us,吞吐量为 748.78Gb/s
下面我们就根据 Softmax 优化那篇文章所提及的点来逐步分析:
优化1 数据Pack
在之前的 高效、易用、可拓展我全都要:OneFlow CUDA Elementwise 模板库的设计优化思路 里很详细的描述了如何做向量化读写,cuda里最大支持 128bit的读写,那么在数据类型为 Float 时,我们即可以将连续的4个 Float 打包到一起,一次性读写,提升吞吐。
有了解过这方面的读者应该就反应过来,诶 CUDA 里 不是刚好有一个类型叫 float4 就是干这件事的么,没错,但是为了更灵活的支持其他数据类型的向量化,我们利用union共享空间的特性实现了一个 Pack 类:
template<typename T, int N>
struct GetPackType {
using type = typename std::aligned_storage<N * sizeof(T), N * sizeof(T)>::type;
};
template<typename T, int N>
using PackType = typename GetPackType<T, N>::type;
template<typename T, int N>
union Pack {
static_assert(sizeof(PackType<T, N>) == sizeof(T) * N, "");
__device__ Pack() {
// do nothing
}
PackType<T, N> storage;
T elem[N];
};
优化2 数据缓存
整个算子逻辑是需要读取一遍数据,计算scale
,然后再读取一遍数据,用scale
进行缩放。很显然这里我们读取了两遍数据,而数据是放在 Global Memory,带宽比较低,会带来读取耗时。
![ac792f49114c3502874e1716dcf7c117.png](https://i-blog.csdnimg.cn/blog_migrate/1a4a7b73e28e713c1eb8b76e402c7927.png)
![ece11bfd8735a863c23afbcb06b90e06.png](https://i-blog.csdnimg.cn/blog_migrate/586def713d61f3478edf36e2dbfe7bb1.png)
一个很自然的想法是缓存到寄存器/Shared Memory中。由于这里我们只实现 WarpReduce 版本,所以我们是缓存到寄存器(其他版本可以参考开头的优化 Softmax 文章)中,减少一次对 Global Memory 的读取。
template<typename T, typename IDX, int pack_size, int cols_per_thread>
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {
// ...
T buf[cols_per_thread];
// ...
优化3 使用Warp处理一行数据
相较 BaseLine,我们这里使用 warp 作为 Reduce 的单位进行操作,首先我们简单看下 WarpReduce 的实现。
template<typename T>
struct AbsMaxOp {
__device__ __forceinline__ T operator()(const T& a, const T& b) const {
return max_func(abs_func(a), abs_func(b));
}
};
template<typename T>
__inline__ __device__ T WarpAbsMaxAllReduce(T val){
for(int lane_mask = kWarpSize/2; lane_mask > 0; lane_mask /= 2){
val = AbsMaxOp<T>()(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
是 0xffffffff
var
则是寄存器值,laneMask
则是用来做按位异或的掩码
这里引入一个概念叫 Lane,它表示线程束中的第几号线程
示意图如下:
![4bb8eefa4c701b3465efc0e93cfef207.png](https://i-blog.csdnimg.cn/blog_migrate/883fae5e3b76c5476769af27b12485fc.png)
当 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<typename T, typename IDX, int pack_size, int cols_per_thread>
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {
// ...
}
跟BaseLine一样,我们block大小还是设置为128个线程,一个warp是32个线程,所以我们一个block可以组织成(32, 4),包含4个warp。
![2348883ac609c2d54b778ca69dd12643.png](https://i-blog.csdnimg.cn/blog_migrate/98e9889e1022c3c2140560c002363468.png)
根据这个层级划分,我们可以计算出:
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<T, pack_size>;
using LoadStorePack = Pack<T, pack_size>;
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 循环,以保证整一行都有被读取到:
![3ad267c261ee933a468d5f25113abc0a.png](https://i-blog.csdnimg.cn/blog_migrate/9c603faf99bf9dc6f494fc13df48f37a.png)
一次性读取到一个 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<LoadStoreType*>(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<T>(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<LoadStoreType*>(x)+ store_offset) = store_pack.storage;
}
完整代码如下:
template<typename T>
__inline__ __device__ T WarpAbsMaxAllReduce(T val){
for(int lane_mask = kWarpSize/2; lane_mask > 0; lane_mask /= 2){
val = AbsMaxOp<T>()(val, __shfl_xor_sync(0xffffffff, val, lane_mask));
}
return val;
}
template<typename T, typename IDX, int pack_size, int cols_per_thread>
__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<T, pack_size>;
using LoadStorePack = Pack<T, pack_size>;
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<LoadStoreType*>(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<T>(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<LoadStoreType*>(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 的使用。
本文仅做学术分享,如有侵权,请联系删文。
干货下载与学习
后台回复:巴塞罗那自治大学课件,即可下载国外大学沉淀数年3D Vison精品课件
后台回复:计算机视觉书籍,即可下载3D视觉领域经典书籍pdf
后台回复:3D视觉课程,即可学习3D视觉领域精品课程
3D视觉精品课程推荐:
2.面向自动驾驶领域的3D点云目标检测全栈学习路线!(单模态+多模态/数据+代码)
3.彻底搞透视觉三维重建:原理剖析、代码讲解、及优化改进
4.国内首个面向工业级实战的点云处理课程
5.激光-视觉-IMU-GPS融合SLAM算法梳理和代码讲解
6.彻底搞懂视觉-惯性SLAM:基于VINS-Fusion正式开课啦
7.彻底搞懂基于LOAM框架的3D激光SLAM: 源码剖析到算法优化
8.彻底剖析室内、室外激光SLAM关键算法原理、代码和实战(cartographer+LOAM +LIO-SAM)
重磅!计算机视觉工坊-学习交流群已成立
扫码添加小助手微信,可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群,旨在交流顶会、顶刊、SCI、EI等写作与投稿事宜。
同时也可申请加入我们的细分方向交流群,目前主要有ORB-SLAM系列源码学习、3D视觉、CV&深度学习、SLAM、三维重建、点云后处理、自动驾驶、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、深度估计、学术交流、求职交流等微信群,请扫描下面微信号加群,备注:”研究方向+学校/公司+昵称“,例如:”3D视觉 + 上海交大 + 静静“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进去相关微信群。原创投稿也请联系。
▲长按加微信群或投稿
▲长按关注公众号
3D视觉从入门到精通知识星球:针对3D视觉领域的视频课程(三维重建系列、三维点云系列、结构光系列、手眼标定、相机标定、激光/视觉SLAM、自动驾驶等)、知识点汇总、入门进阶学习路线、最新paper分享、疑问解答五个方面进行深耕,更有各类大厂的算法工程人员进行技术指导。与此同时,星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息,打造成集技术与就业为一体的铁杆粉丝聚集区,近4000星球成员为创造更好的AI世界共同进步,知识星球入口:
学习3D视觉核心技术,扫描查看介绍,3天内无条件退款
圈里有高质量教程资料、可答疑解惑、助你高效解决问题
觉得有用,麻烦给个赞和在看~