CUDA 算子优化：Softmax

1.softmax定义

把输入映射为0-1之间的实数，并且归一化保证和为1。max非黑即白，soft，输出是每个分类被取到的概率。

定义：

该元素的softmax值就是该元素的指数和所有元素指数和的比值

定义的原因：

1.softmax设计的初衷是希望特征对概率的影响是乘性的。

2.多类目标分类问题的目标函数常常选为cross-entropy。

2.softmax求导

回顾交叉熵损失函数cross entropy

神经网络最后一层的计算方式如下：

1.网络的最后一层得到每个类别的scores。

2.score与sigmoid函数或者softmax函数计算得到概率输出

3.类别概率与真实类别的one-hot形式进行交叉熵计算。

yi表示类别是1，^yi表示预测类别为1的概率。

在多分类问题中，经常使用交叉熵作为损失函数与softmax配合使用

$LOSS=-\sum_{i}^{}t_{i}lny_{i}$

softmax函数

$y_{i} =\frac{e^{i}}{\sum_{j}^{}e^{j}}=1-\frac{\sum_{j\neq i}^{}e^{j}}{\sum_{j}^{}e^{j}}$

对LOSS求导=yi-1得到损失函数

所以，只需要正向求出yi，将结果-1就是反向更新的梯度。

3.如何实现一个高效的softmax CUDA kernel

Kernel：GPU的基本计算任务描述单元。每个Kernel都会根据配置参数在GPU上由非常多个线程并行执行，GPU计算高效就是因为同时可以由i数千个core（thread）同时执行

GPU的线程在逻辑上分为Thread/Block/Grid三级，硬件上分为core/warp两级

GPU内存分为Global memory/Shared memory/Local memory三级

GPU最主要提供的是两种资源：计算资源和显存带宽资源，优化目标是尽可能充分利用显存带宽资源

4.如何评估一个CUDA Kernel是否充分利用了显存带宽资源？

充分利用指的是Kernel的有效显存读写带宽达到了设备显存带宽的上限，其中设备显存带宽可以通过执行cuda中的bandwidthTest得到。Kernel的有效显存带宽通过Kernel读写数据量和Kernel执行时间进行评估：

当前Kernel的有效显存带宽 = 读写数据量/执行时间

Naive的softmax实现：未经优化的softmax Kernel的最高理论带宽。一个基础的softmax计算中，分别调用几个基础的CUDA Kernel function来完成整体的计算：

假设输入的数据大小为D，shape=(num_rows,num_cols)，即D=num_rows * num_cols，最naive 的操作会多次访问Global memory，其中：

ReduceMax = D+num_rows(read为D，write为num_rows)
BroadcaseSub = 2*D + num_rows(read为D+num_rows, write为D)
Exp = 2*D（read/write均为D）
ReduceSum=D+num_rows(read为D，write为num_rows)
BroadcastDive = 2*D+num_rows(read为D+num_rows，write为D)

总共需要8*D+4*num_rows的访存开销。num_rows相比于D可以忽略，则Navie版本的softmax CUDA kernel需要访问至少8倍数据的显存，即：

NaiveSoftmaxKernel有效显存带宽<理论带宽/8

对于3090显卡，其理论带宽上限为936GB/S，Navie版本的Softmax CUDA Kernel利用显存带宽的上界就是936/8=117GB/S。

softmax的shape=（49152，num_cols），其中49152=32*12*128，BERT-base网络中的attention Tensor的前三维，固定前三维，最后一维动态变化

5.OneFlow深度优化Softmax CUDA Kernel的技巧

Softmax函数的输入形状：（num_rows,num_cols）,num_cols的变化，会对有效带宽产生影响；没有一种通用的优化方法可以实现在所有的num_cols的情况下都是传输最优的。所以，在OneFlow中采用分段函数优化SoftmaxKernel：对于不同num_cols范围，选择不同的实现，以在所有情况下都能达到较高的有效带宽。

oneFlow中分三种实现，分段对softmax进行优化：

（1）一个Warp处理一行的计算，适用于num_cols<=1024情况

硬件上并行执行的32个线程称为一个warp，同一个warp的32个thread执行同一条指令。warp是GPU调度执行的基本单元。

（2）一个Block处理一行的计算，借助Shared Memory保存中间的结果数据，适用于需要的Shared Memory资源满足Kernel Launch的可启动条件的情况，在本测试环境中是1024<num_cols<=4096

（3）一个Block处理一行的计算，不使用Shared Memory，重复读输入X，适用于不支持（1）（2）的情况。

实现1：每个Warp处理一行元素

每个Warp处理一行元素，行Reduce需要做Warp内的Reduce操作，这里实现WarpAllReduce完成Warp内各线程间的求Global Max和Global Sum操作，WarpAllReduce是利用Warp级别原语__shfl_xor_sync实现的，代码如下。

//声明一个模板函数WarpReduce，接受他的两个模板函数：
//1.ReductionOp：一个模板参数，表示一个接受单个类型参数的模板，该模板应产生一个可调用的对象，用于定义如何归约两个值。
//2.T：要归约的数据类型
template<template<typename> typename ReductionOp, typename T>

//声明一个内联设备函数WarpReduce，该函数在GPU上执行，并接受一个类型为T的val，返回类型为T
__inline__device__ T WarpAllReduce(T val) {
//初始化mask为warp的一半，oxffffffff用于指定那些线程参与同步，全部参与
//当mask>0时，循环继续。这个循环用于在不同层次上将warp中的值进行归约
//利用__shfl_xor_sync将当前线程值与warp中基于当前mask的其他线程的值进行交换
//ReductionOp<T>()(val, ...)调用归约操作来合并当前线程的值和shuffle得到的值
//mask/=2：mask减半，准备进行下一轮的shuffle和归约操作
    for (int mask = kWarpSize / 2; mask > 0; mask / = 2) {
        val = ReductionOP<T>()(val, __shfl_xor_sync(oxffffffff, val,mask));
    }
//返回归约后的值
    return val;
}

主体循环逻辑，首先根据num_cols信息算出每个线程要处理的cols_per_thread，每个线程分配cols_per_thread大小的寄存器，将输入x读到寄存器中，后续计算均从寄存器中读取。

理论上来说，每个warp处理一行元素的速度是最快的，但是由于需要使用寄存器将输入x缓存起来，而寄存器资源是有限的，并且在num_col>1024情况下，使用shared memory 方法已经足够快乐，因此仅在num_cols<=1024时采用Warp的实现。

//pack_size:每个线程处理的数据包大小
//cols_per_thread:每个线程处理的列数
//padding：是否使用填充（输入矩阵是否已经被填充到warpSize的倍数）
template<typename T, int pack_size, int cols_per_thread, bool padding>
__global__ void SoftmaxWarpImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
//确保每个线程处理的列数时每个线程处理的数据包的大小的倍数
  static_assert(cols_per_thread % pack_size == 0, "");
//每个线程处理的包大小
  constexpr int num_packs = cols_per_thread / pack_size;
  //列数不超过每个warp处理的列数
  assert(cols <= cols_per_thread * kWarpSize);
  //为每个线程分配一个缓冲区存储数据
  using ComputeType = typename GetComputeType<T>::type;
  ComputeType buf[cols_per_thread];
  //计算全局warp的id
  const int global_warp_id = blockIdx.x * blockDim.y + threadIdx.y;
  //计算全局warp的数量
  const int num_global_warp = gridDim.x * blockDim.y;
  //线程id
  const int lane_id = threadIdx.x;
  //遍历行
  for (int64_t row = global_warp_id; row < rows; row += num_global_warp) {
    //计算当前行在输入输出的偏移量
    const int64_t row_offset = row * cols;
    //计算当前行在输入输出的指针
    const T* row_x = x + row_offset;
    T* row_y = y + row_offset;
    //计算当前行最大的线程数量
    ComputeType thread_max = -Inf<ComputeType>();
#pragma unroll//表示展开循环
//循环处理每个包
    for (int pack_id = 0; pack_id < num_packs; ++pack_id) {
    //计算当前线程处理的列索引
      const int col = (pack_id * kWarpSize + lane_id) * pack_size;
      //如果不使用填充或者列在有效范围内
      if (!padding || col < cols) {
        //将输入数据加入到缓冲区
        MultiFetch<T, ComputeType, pack_size>()(buf + pack_id * pack_size, row_x + col);
#pragma unroll
        //找出线程中的最大值
        for (int i = 0; i < pack_size; ++i) {
          thread_max = max(thread_max, buf[pack_id * pack_size + i]);
        }
      } else {
        //否则，设为负无穷，初始化buf数组中的一部分将其设为无穷大，确保softmax计算中0或者非常小的数字不会被忽略
#pragma unroll
        for (int i = 0; i < pack_size; ++i) { buf[pack_id * pack_size + i] = -Inf<ComputeType>(); }
      }
    }

    //计算warp（一组并行的线程）的最大值，MaxOp定义了如何比较和选择最大值的操作或函数对象
    const ComputeType warp_max = WarpAllReduce<MaxOp, ComputeType>(thread_max);
    ComputeType thread_sum = 0;

#pragma unroll
    //计算softmax并求和，为了计算稳定性，从每个值中减去warp最大值然后取指数
    for (int i = 0; i < cols_per_thread; ++i) {
      buf[i] = exp(buf[i] - warp_max);
      thread_sum += buf[i];
    }

    //计算warp内的和，计算warp内所有线程的和
    const ComputeType warp_sum = WarpAllReduce<SumOp, ComputeType>(thread_sum);
#pragma unroll
    //归一化，使和为1，呈概率分布形式
    for (int i = 0; i < cols_per_thread; ++i) { buf[i] = buf[i] / warp_sum; }
#pragma unroll
    //将结果存回数组或者原始矩阵
    for (int i = 0; i < num_packs; ++i) {
      const int col = (i * kWarpSize + lane_id) * pack_size;
      if (!padding || col < cols) {
        MultiStore<ComputeType, T, pack_size>()(row_y + col, buf + i * pack_size);
      }
    }
  }
}

实现2：一个Block处理一行元素

一个Block处理一行元素，行Reduce需要做Block内的Reduce操作，需要做Block内线程同步，利用BlockAllReduce完成Warp内各线程间的求Global Max和Global Sum操作。BlockAllReduce时借助Cub的BlockReduce方法实现的

//ReductionOp：模板参数，T：用于归约操作的数据类型，block_size：CUDA块大小
template<template<typename> typename ReductionOp, typename T, int block_size>
__inline__ __device__ T BlockAllReduce(T val) {
  typedef cub::BlockReduce<T, block_size> BlockReduce;
  //声明共享内存变量，用于在块内的归约操作中存储中间结果
  __shared__ typename BlockReduce::TempStorage temp_storage;
  //声明  共享内存变量result_broadcast，用于在块内广播归约操作的结果
  __shared__ T result_broadcast;
  //执行归约操作
  T result = BlockReduce(temp_storage).Reduce(val, ReductionOp<T>());
  if (threadIdx.x == 0) { result_broadcast = result; }
  //线程同步
  __syncthreads();
  return result_broadcast;
}

主体循环逻辑，根据num_cols算出需要的Shared Memory大小作为Launch Kernel参数，借助Shared Memory保存输入，后续计算直接从Shared Memory读取。

由于SM内的Shared Memory资源同样有限，因此当num_cols超过一定范围，kernel启动时申请Shared Memory超过最大限制，就会出现无法启动的问题，因此，仅在调用cudaOccupancyMaxActiveBlocksPerMultiprocessor返回值大于0时采用Shared Memory方案。

此外，由于block内线程要做同步，当SM中正在调度执行一个Block到达同步点时，SM内可执行Warp逐渐减少，若同时执行的Block只有一个，则SM中可同时执行的Warp会在此时逐渐降为0，会导致计算资源空闲，造成浪费，若此时同时有其他Block在执行，则在一个Block到达同步点时仍然有其他Block可以执行。当block_size越小时，SM可以同时调度的Block越多，因此在这种情况下block_size越小越好。但是当在调大block_size，SM能同时调度的Block数不变的情况下，block_size越大越好，越大就有越好的并行度。因此代码中在选择block_size时，对不同block_size都计算了cudaOccupancyMaxActiveBlocksPerMultiprocessor，若结果相同，使用较大的block_size。

template<typename T, int pack_size, int block_size>
__global__ void SoftmaxBlockSMemImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
  using ComputeType = typename GetComputeType<T>::type;
  //共享内存定义
  extern __shared__ __align__(sizeof(ComputeType)) unsigned char shared_buf[];
  auto* buf = reinterpret_cast<ComputeType*>(shared_buf);
  //获取当前线程id
  const int tid = threadIdx.x;
  assert(cols % pack_size == 0);
  const int num_packs = cols / pack_size;
  //外层循环用来迭代每一行。由于每一个线程块可能处理多行数据，使用blockidx.x确定当前线程块处理的起始行
  for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
    //计算当前行的偏移量，并初始化输入和输出指针。初始化一个变量来跟踪线程的最大值
    const int64_t row_offset = row * cols;
    const T* row_x = x + row_offset;
    T* row_y = y + row_offset;
    ComputeType thread_max = -Inf<ComputeType>();
    //数据包加载和最大值查找
    for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
      ComputeType pack[pack_size];
      MultiFetch<T, ComputeType, pack_size>()(pack, row_x + pack_id * pack_size);
#pragma unroll
      for (int i = 0; i < pack_size; ++i) {
        buf[i * num_packs + pack_id] = pack[i];
        thread_max = max(thread_max, pack[i]);
      }
    }
    //使用块内归约操作找到整个块的最大值和总和
    const ComputeType row_max = BlockAllReduce<MaxOp, ComputeType, block_size>(thread_max);
    ComputeType thread_sum = 0;
    for (int col = tid; col < cols; col += block_size) {
      const ComputeType exp_x = exp(buf[col] - row_max);
      buf[col] = exp_x;
      thread_sum += exp_x;
    }
    const ComputeType row_sum = BlockAllReduce<SumOp, ComputeType, block_size>(thread_sum);
    //计算softmax值并存储
    for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
      ComputeType pack[pack_size];
#pragma unroll
      for (int i = 0; i < pack_size; ++i) {
        pack[i] = buf[i * num_packs + pack_id] / row_sum;
        thread_max = max(thread_max, pack[i]);
      }
      MultiStore<ComputeType, T, pack_size>()(row_y + pack_id * pack_size, pack);
    }
  }
}

实现3：一个Block处理一行的元素，不是用Shared Memory

重复读输入x和实现2一样，仍然是一个Block处理一行元素，不同的是，不再用Shared Memory缓存x，而是在每次计算时重新读输入x，这种实现没有最大num_cols的限制，可以支持任意大小。

此外需要注意的是，在这种实现中，block_size应该设越大越好，block_size越大，SM中能同时并执行的Block数就越少，对cache的需求就越少，就有更多机会命中Cache，多次读x不会多次访问Global Memory，因此在实际测试中，在能利用Cache情况下，有效带宽不会因为读3次x而降低几倍。

//定义一个CUDA核函数模板，接受三个模板参数：数据类型T，数据包大小pack_size，线程块大小
template<typename T, int pack_size, int block_size>
__global__ void SoftmaxBlockUncachedImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
  //使用computeType模板元函数来确定用于计算的数据类型
  using ComputeType = typename GetComputeType<T>::type;
  //获取当前线程ID
  const int tid = threadIdx.x;
  //断言：确保列数是数据包大小的整数倍
  assert(cols % pack_size == 0);
  //计算数据包数量
  const int num_packs = cols / pack_size;
  //外层循环 迭代每一行
  for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
    //计算当前行的偏移量
    const int64_t row_offset = row * cols;
    //指向当前行的输入和输出数据的指针
    const T* row_x = x + row_offset;
    T* row_y = y + row_offset;
    //初始化线程最大值为负无穷大
    ComputeType thread_max = -Inf<ComputeType>();
    //内层循环：加载数据包并找到整个线程块的最大值
    for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
      ComputeType pack[pack_size];
      MultiFetch<T, ComputeType, pack_size>()(pack, row_x + pack_id * pack_size);
      //展开循环，找到当前线程的最大值
#pragma unroll
      for (int i = 0; i < pack_size; ++i) { thread_max = max(thread_max, pack[i]); }
    }

    //使用blockallreduce模板函数找到整个线程块的最大值
    const ComputeType row_max = BlockAllReduce<MaxOp, ComputeType, block_size>(thread_max);
    ComputeType thread_sum = 0;
    //再次遍历数据包，计算每个数据包的指数和
    for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
      ComputeType pack[pack_size];
      MultiFetch<T, ComputeType, pack_size>()(pack, row_x + pack_id * pack_size);
      //展开循环，计算指数和
#pragma unroll
      for (int i = 0; i < pack_size; ++i) { thread_sum += exp(pack[i] - row_max); }
    }
    //使用blockallreduce模板函数找到整个线程块的和
    const ComputeType row_sum = BlockAllReduce<SumOp, ComputeType, block_size>(thread_sum);
    for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
      ComputeType pack[pack_size];
      MultiFetch<T, ComputeType, pack_size>()(pack, row_x + pack_id * pack_size);
      //展开循环，计算softmax值
#pragma unroll
      for (int i = 0; i < pack_size; ++i) { pack[i] = exp(pack[i] - row_max) / row_sum; }
      //使用multistore模板函数将结果存储到全局内存中
      MultiStore<ComputeType, T, pack_size>()(row_y + pack_id * pack_size, pack);
    }
  }
}

6.公共优化技巧

1.将Half类型pack成Half2进行存取，在不改变延迟情况下提高指令吞吐，类似CUDA template for element-wise kernels 的优化

2.Shared Memory中的 Bank Conflicts

CUDA将Shared Memory按照4字节或8字节大小划分到32个bank中，对于Volta架构是4字节，以4字节为例，0-128bytes地址分别在bank0-31中，128-256也分别在bank0-31中。

当一个Warp内的不同线程访问同一个bank 的不同地址，就会出现Bank Conflicts。当发生bank冲突时，线程间只能顺序访问，增加延迟，

若一个Warp内每个线程读4个字节，顺序访问，则不会有Bank Conflicts，若一个Warp内每个线程读8个字节，则Warp内0号线程访问的地址在第0和第1个bank，1号线程访问的地址在第2和第3个bank，以此类推，16号线程访问地址又在第0和第1个bank内，和0号线程访问了同一个bank的不同地址，此时即产生了bank冲突

在实现2中，给Shared memory赋值过程中，若采用下面方法，当pack_size=2，每个线程写连续两个4byte地址，就会产生bank conflicts

#pragma unroll
  for(int j = 0; j < pack_size; ++j) {
    buf[pack_id * pack_size * j] = pack[j];
    thread_max = max(thread_max, pack[j]);
  }

因此，在实现（2）中，对Shared memory采用了新的内存布局，避免了同一个Warp访问相同bank的不同地址，避免了Bank Conflicts。

#pragma unroll
  for(int j = 0; j < pack_size; ++j) {
    buf[num_packs * j + pack_id] = pack[j];
    thread_max = max(thread_max, pack[j]);
  }