5-GGML:看一下CUDA加法算子!

GPU加法

主要这里看一下CUDA的加速算法是怎么写的了,CUDA的并行基本语法规则这里就不再科普了,可以去查阅一下其他的资料什么的.
代码比较难找,可以直接搜这函数bin_bcast_cuda

更改维度

CUDA的block块和grad块的维度都是三维的,但是GGML的Tensor维度是4维度,此时简单的理解就是把两个维度合并成一个就好了.那么思考一下怎么合并合理,首先肯定是挨着的合并更好,其次如果我矩阵是有序存储的,合并之后我也希望是有序的,该怎么合,显然应该合入第0维度和第1维度.

auto collapse = [](int64_t cne[]) {// 合并0 1维cne[0] *= cne[1];// 合并以后0维度变成 0 * 1个cne[1] = cne[2];cne[2] = cne[3];cne[3] = 1;};auto collapse_nb = [](size_t cnb[], const int64_t cne[]) {cnb[1] *= cne[1];// 扩展自然会变大cnb[2] *= cne[2];cnb[3] *= cne[3];};// 如果是连续存储的可以进行维度合并操作if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)) {for (int i = 0; i < 4; i++) {if (nr[i] != 1) {// 需要扩展的话不进行运算break;}if (i > 0) {collapse_nb(cnb, cne);collapse_nb(cnb0, cne0);collapse_nb(cnb1, cne1);collapse(cne);collapse(cne0);collapse(cne1);}}}

block_dims与block_nums设置

这里注意的问题,

1. x轴和ne0相关,y轴和ne1相关,z轴和ne2和ne3相关
2. x轴并没有设置满(我这里满的意思是x轴并不是全部的,而是he0一半这么多元素)其他轴是满的
3. z轴是包含ne2和ne3的,后续可根据2\3的维度信息进行还原

const int block_size = 128;int64_t hne0 = std::max(ne0/2LL, 1LL);// 取ne0的一半作为x维度dim3 block_dims;
// 输出张量第0维度的一半,至少为1,最大为min(block_size, ne0/2)
block_dims.x = std::min<unsigned int>(hne0, block_size);
// 输出张量第1维度,最大为 ne1 维度 block_size / block_dims.x
block_dims.y = std::min<unsigned int>(ne1, block_size / block_dims.x);
// 最大是64
block_dims.z = std::min(std::min<unsigned int>(ne2*ne3, block_size / block_dims.x / block_dims.y), 64U);dim3 block_nums((hne0 + block_dims.x - 1) / block_dims.x,// 0维度差的(ne1 + block_dims.y - 1) / block_dims.y,// 1维度差的(ne2*ne3 + block_dims.z - 1) / block_dims.z// 2维度和3维度合并后的差的);

调用CUDA的核函数

这里直接看else吧,就是调用CUDA的东西,我们看一下核函数内部的实现(以k_bin_bcast为例子).

if (block_nums.z > 65535) {// this is the maximum number of blocks in z dimension, fallback to 1D grid kernelint block_num = (ne0*ne1*ne2*ne3 + block_size - 1) / block_size;k_bin_bcast_unravel<bin_op><<<block_num, block_size, 0, stream>>>(src0_dd, src1_dd, dst_dd,ne0, ne1, ne2, ne3,ne10, ne11, ne12, ne13,/* s0, */ s1, s2, s3,/* s00, */ s01, s02, s03,/* s10, */ s11, s12, s13);
} else {k_bin_bcast<bin_op><<<block_nums, block_dims, 0, stream>>>(src0_dd, src1_dd, dst_dd,ne0, ne1, ne2, ne3,ne10, ne11, ne12, ne13,/* s0, */ s1, s2, s3,/* s00, */ s01, s02, s03,/* s10, */ s11, s12, s13);
}

k_bin_bcast函数

首先看一下i2\i3,以上分析了i2\i3和z轴关联的,所以用z轴的索引加3维度就可区分开了.然后i1是不用说了只和y轴相关,i0s只和x轴相关.然后可以根据索引把src1矩阵的索引确定下来,注意的是src1是在这里进行的求余操作是进行的缩放(之前加法分析过的允许两个矩阵非完全一致,只要是倍数关系就可以).
详细说一下 "for (int i0 = i0s; i0 < ne0; i0 += blockDim.xgridDim.x) " 这句话,可以看到循环是i0s这个块的开始行执行循环,然后结尾是ne0(也就是最大的行数),循环步长是blockDim.xgridDim.x,为什么循环步长设置成这个呢?因为在block_dims与block_nums设置的时候是没给x维度设置满的,所以在这里进行了循环补偿.比如矩阵0维度有600行,但是0维度设置时的规则是block_dims.x = min(128, 600/2)为128,然后block_nums.x = (600/2 + 128)/128, 本质是有缺失的,需要这样进行补全运算.

template<float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename dst_t>
static __global__ void k_bin_bcast(const src0_t * src0, const src1_t * src1, dst_t * dst,int ne0, int ne1, int ne2, int ne3,int ne10, int ne11, int ne12, int ne13,/*int s0, */ int s1,  int s2,  int s3,/*int s00,*/ int s01, int s02, int s03,/*int s10,*/ int s11, int s12, int s13) {const int i0s = blockDim.x*blockIdx.x + threadIdx.x;const int i1 = (blockDim.y*blockIdx.y + threadIdx.y);const int i2 = (blockDim.z*blockIdx.z + threadIdx.z) / ne3;const int i3 = (blockDim.z*blockIdx.z + threadIdx.z) % ne3;if (i0s >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) {return;}// 扩张的部分const int i11 = i1 % ne11;const int i12 = i2 % ne12;const int i13 = i3 % ne13;// 计算内存偏移const size_t i_src0 =  i3*s03 +  i2*s02 +  i1*s01;const size_t i_src1 = i13*s13 + i12*s12 + i11*s11;const size_t i_dst  =  i3*s3  +  i2*s2  +  i1*s1;// 计算行位于的地址const src0_t * src0_row = src0 + i_src0;const src1_t * src1_row = src1 + i_src1;dst_t * dst_row = dst + i_dst;// 进行计算for (int i0 = i0s; i0 < ne0; i0 += blockDim.x*gridDim.x) {const int i10 = i0 % ne10;dst_row[i0] = (dst_t)bin_op(src0 ? (float)src0_row[i0] : 0.0f, (float)src1_row[i10]);}
}