pytorch算子的cuda实现——expand,where,softmax

Answer 1

算子的简单介绍

orch.Tensor.expand() 是 PyTorch 中 Tensor 类的一个方法，用于扩展张量的维度。输入：input 是要扩展的张量，size 是一个元组，指定了要扩展的每个维度的大小。输出：返回一个新的张量，形状是 input 张量的形状扩展后的形状。例如，x 的形状假设为(2,1,5,1），x.expand(2,4,5,6) 表示希望将 x 扩展为(2,4,5,6）形状的数组，注意 expand 只能在原始形状扩张，并且是原始维度的长度为1的地方。

torch.where() 是 PyTorch 中的一个函数，用于根据给定的条件从两个张量中选择元素。输入：condition, inputx, inputy。注意 condition, inputx, inputy 不要求形状保持一致。例如，condition=[2,1,1,1]，inputx=[1,3,4,1]，inputy=[1,2,4,2]，这都符合要求，此时输出的形状应该是[2,3,4,2]，换句话说，where 应该支持广播运算。

一维情况的 expand 和 where

对于一维向量来说，使用 expand 和 where 实现非常简单。对于 expand，我们可以写出伪代码。对于 where，我们也可以写出伪代码。

高维数组的 expand

对于高维数组而言，expand 需要考虑在哪个维度扩张。假设输入向量 input 的形状为 [公式]，其中至少存在一个 [公式]（这是因为 expand 只能在某个维度的长度为1的地方扩张）。假设经过 expand 输出以后 output 的形状为 [公式]。由于计算机存储高维数组仍然需要将数组拉平，弄成一维数组存储，因此我们需要获得高维数组拉平后的长度以及全局索引。例如，input[inputIdx]: inputIdx =[公式]。output[outputIdx]: outputIdx =[公式]。由此我们就可以直接获得 expand 的公式。假设 [公式]，即扩张的维度 dim=1，此时对于 output 的任意一个元素，假设索引为 [公式]，那么这个元素对应的取值就应该是 input 中索引为 [公式] 的元素。

高维数组的 where 编程

有了 expand 的经验之后，where 的实现就相对容易理解了。在编写 expand 算子的过程中，核心内容是根据 outputIdx 反推得到 inputIdx。同样地，我们可以根据 outputIdx 计算出 inputxIdx, inputyIdx, conditionIdx，然后直接计算即可。

一维向量的 softmax 介绍

softmax 的实现相对复杂。这里本人暂时只能实现一维向量的 softmax。假设输入是 x，输出是 y，其中 [公式]。在处理过程中，为了避免指数变化导致溢出，首先需要计算出向量 x 的最大值 [公式]，然后做一次变换 [公式]，将变换以后的数组记录为 [公式]，从而得到 [公式]。

softmax 的实现可以分解为三步：

遍历向量 x，计算得到 M，也就是一个 max 规约操作。
更新向量 x，其中 [公式]，计算 [公式]，也就是一个 sum 规约。
得到向量 y，其中 [公式]。

向量的 sum 规约

softmax 涉及到两次规约，这里我们先介绍使用的规约策略。我们先用 sum 规约介绍，其实 max 规约几乎如出一辙。

交叉配对规约：reduce=0

向量长度是 size（假设 size 是 2 的幂次方），交叉配对的原理是不断将奇数位置的元素加到偶数位置，最终将所有元素的和存放在起始位置。

以代码举例说明，这里重点说明第 65 行的代码。定义了一个数组 tmp，其中 __shared__ 关键字表示这个数组放在共享内存中，以及 tmp 数组面对当前线程块中的所有线程都是共享的。对于每个线程块来说，tmp 存储的是该线程块每个线程对应的数据。__syncthreads() 在 GPU 端做一个同步，确保所有线程的数据都赋值到了 tmp。重点来看 65 行以后的代码。

以图为例，blockDim=(8,1,1), gridDim=(4,1,1)，一共有 32 个线程，blockDim.x=8。考虑第 65 行这个循环。当 strip=1 时，threadIdx.x%2 == 0，则 threadIdx.x 有 0，2，4，6 满足条件。当 threadIdx.x=0 时，67 行代码就是 tmp[0] += tmp[1]，类似地，成功将数据从 8 个位置转移到了 4 个位置，然后继续修改 strip=2,4，最后将所有元素的和转移到 tmp[0] 即可。最后将 tmp[0] 的数据放到 result[blockIdx.x] 中，也就是说 result 的每个元素存储的都是对应某个线程块内所有元素的和。最终结果全在 result 数组中，那么我们就可以推断出 result 的长度应该是 ceil(num_rows / BLOCK_DIM)，其中 ceil 表示向上取整，num_rows 是向量的长度。

交叉配对规约的特点是规约过程类似于一个二叉树，每次都把奇数位置的元素添加到偶数位置，架构很简单，但具体每一次规约的数据访问却是跳跃的。

交错配对规约：reduce=1

分析后面的规约代码，当 strip=4 时，threadIdx.x < 4，就有 threadIdx.x=0,1,2,3 满足条件。当 threadIdx.x=0 时，就有 tmp[0] += tmp[4]，tmp[1] += tmp[5]，tmp[2] += tmp[6]，tmp[3] += tmp[7]。这种做法的好处是每次规约数组的访问都是连续的，可以减少访存开销。后面的数值实验也证明这种并行策略确实加速效果更好。

shuffle warp：reduce=2

这种做法的原理比较复杂，这里本人解释的不一定准确。首先声明了一个长度为 32 的共享内存数组 tmp，这个 32 实际上是线程的 warp，GPU 线程的设计是每 32 个线程为一组，称为一个线程簇 warp，我们的一维线程块的 BLOCK_DIM 每隔 32 个线程做一个切分。

现在我们进一步考虑线程块内部的线程，对于某个线程块里面的某一个 warp 而言，这里面的 data 其实是线程私有变量，但是 data += __shfl_down_sync(0xffffffff, data, 16) 的意思表示不同线程之间的 data 其实可以通过某种特殊方式相互读取累加的。该 warp 有 32 个线程，第一次累加的原则是，第 0 个线程和第 16 个线程相加，...，第 15 个线程和第 31 个线程数据相加，这样就把 warp 内 32 个线程的数据全部累加到了前 16 个线程。后面几句代码也是类似的，最后把该 warp 的 32 个线程数据累加到第一个线程。所以当 threadIdx.x%32==0 时，表示 threadIdx.x 这个线程存储了该 warp 中 32 个线程数据的累加和。所以 tmp[threadIdx.x/32] 位置存储的就是每一个 warp 所有线程数据的和。

现在需要把 tmp 这 32 个数据类似原则累加，因此借助于 threadIdx.x 可以直接读取 tmp 的数据，但是 threadIdx.x 的范围太大，我们不要借助那么多所以需要做一个判断以免访问 tmp 越界。最后将 tmp 的数据做完规约以后得到的 data 就保存了当前线程块的所有线程数据的累加和。而我们线程块的总数是 gridDim.x，blockIdx.x 的范围就是 0 到 gridDim.x-1，所以 result 的长度还是 gridDim.x，每个元素保存了对应线程块所有数据的累加和。

对于某个 warp，为什么第一次规约是以 16 为距离累加呢？实际上，从线程的角度来看，比如我们考虑最前面的 32 个线程，也就是 block 中第一个 warp，0 号线程一直到 15 号线程做累加的时候读取的数据正好是相邻的 16 个元素，累加的时候读取的是后面 16 个元素。不管怎么累加，不同线程之间读取的数据总数相邻，这样就保证了数据的局部性，所以访问时间进一步缩短。

把这些并行策略封装在 pre_dot 以后，后面仍然是写一个 cuda_dot 函数将 result 的结果传回 CPU 的 host_z，然后再串行求 host_z 的元素和。特别注意的是此时调用需要使用 pre_dot<<>>(cuda_x, cuda_y, cuda_z, result)，share_size 表示申请的共享内存大小，这里我一般使用 BLOCK_DIM×sizeof(double)。

向量的 max 规约

前面两种规约方式对于 max 来说是一样的，稍微有一点区别的是 reduce=2 这种抽牌规约策略。有兴趣的同学可以查看一下 __shfl_down_sync 函数的用法就知道这个是怎么回事了。

初级 softmax

有了上面的 sum 和 max 规约策略，我们可以给一个初级的 softmax 算法，步骤如下：

1：先对数组 input 做一个 max 规约，此时会得到一个 result 数组，这个数组的长度是 ceil(N/BLOCK_DIM)，其中 N 是数组的长度，BLOCK_DIM 是线程块的大小。观察向量的 max 规约可以发现这个规约其实只是对向量做了切割，切割成了 ceil(N/BLOCK_DIM)，每一份长度是 BLOCK_DIM。而这个向量 result 的长度正好是 ceil(N/BLOCK_DIM)，其中 result 的每个元素存储的是对应的那份长度为 BLOCK_DIM 的数据的 max，但是这个不是全局 max，因此我们需要把向量 result 从 GPU 搬运到 CPU，然后通过一个串行 for 循环来遍历 result，由此来获得全局 max，假设是 M。

2：然后我们做一个更新向量 x，其中 [公式]，计算 [公式]，此时这个做 sum 规约，仍然是把向量切割成了 ceil(N/BLOCK_DIM)，每一份长度是 BLOCK_DIM，并且每份长度为 BLOCK_DIM 的数据的 sum 被存储到了 result 的对应位置（result 可以多次利用，不需要重新定义新的数组）。然后还是需要把向量 result 从 GPU 搬运到 CPU，然后通过一个串行 for 循环来遍历 result，由此来获得全局 sum，假设是 S。

3：得到向量 y，其中 [公式]。这个过程比较无聊，不赘述。

综合上面信息，我们发现这个初级的 softmax 的数据来来回回从 CPU 和 GPU 移动了好几次，效率不太高。

高级 softmax

我们回顾一下上面的步骤 1 和 2，我们发现生成的 result 需要从 GPU 移动到 CPU，然后还要串行遍历 result 计算全局结果，但我们可以类似在 GPU 端规约 result，依次来减少数据移动，增大速度。经过本人测试，这个想法完全可行，可以称之为二次规约：

观察上面这段代码，本人定义了一个 girdmax 函数来实现 result 的规约，这个过程和 reduce=1 的规约策略一模一样，经过这样规约以后，全局的 max 将会被存储到 result[0]，也就是 result 向量的开头元素。

类似地，我们继续这么做，这个 gridsum 会将 result 的所有元素都加载到 index=0，也就是开头位置，这样就得到了我们的高级 softmax。

高维数组的 softmax 目前还没调节好代码，有兴趣的同学可以自己试一试，完整代码参考本人博客。

本人在下一篇博客给出了一个 softmax 算子融合的介绍，有兴趣可以参考。