GPU CUDA进行高性能reduce计算的几种方法

Luchang-Li

已于 2024-12-18 12:32:55 修改

阅读量1k

点赞数 16

分类专栏：深度学习编译器推理引擎文章标签： cuda reduce

于 2024-09-11 16:52:01 首次发布

本文链接：https://blog.youkuaiyun.com/u013701860/article/details/142144398

版权

推理引擎同时被 2 个专栏收录

33 篇文章

订阅专栏

深度学习编译器

19 篇文章

订阅专栏

背景

Reduce在深度学习中是一种非常常见的计算类型。例如求取张量整体或者某个维度的最大、最小值、均值、求和等的专一的Reduce算子；此外，在Softmax以及LayerNorm, RmsNorm，GroupNorm等各种Normalization算子中需要求均值和标准差，这也依赖于reduce计算；另外对于向量矩阵乘的场景也通常采用reduce计算模式。因此如何利用GPU多线程进行高效协作进行reduce计算非常重要。本文讨论了如何基于GPU开发高性能的reduce计算方法。

Reduce任务类型

Reduce任务一般是把一个张量沿着某一个维度进行reduce计算，而这可以简单分类为两种类型：对最内层维度（axis=-1）进行reduce，和对非最内层维度进行reduce。如下图所示，如果对一个shape为[3,4]的张量进行reduce，沿着axis=-1进行reduce得到一个shape为[3]的张量，而沿着axis=0进行reduce，得到一个shape为[4]的张量。对于维度大于2D的张量，可以把没有进行reduce的维度进行合并。如果对多个连续的axes reduce，同样可以把这些连续的axes合并为一个axis，从而大多数场景的任务类型都可以转换到前面所述的两种任务类型。

Reduce线程协作方式

在单个GPU中对一个张量最内层或某个中间维度进行reduce计算的多线程任务协作方式上，根据张量大小可以大体分为：

1. thread reduce,：单个线程独立完成整个维度的reduce，无需与其他线程协作。

2. sub-warp reduce：采用少量的多个线程进行协作完成reduce，但是这里针对任务量比较小的场景，通常使用4个或者8个线程进行协作。这种方式作者独立提出，而其他博客很少提到。这种协作的粒度比下面采用一个warp的32个线程协作更小。

3. warp reduce：一个warp的32个线程协作完成reduce。

4. thead block reduce：一个线程块的所有线程协作完成，而一个线程块通常包括多个warp，例如达到1024个线程。

5. multiple thead block reduce: 多个线程块协作完成。

6. grid reduce: 整个grid完成一个reduce计算后得到一个结果，这个在深度学习任务中非常少见。

本文主要针对深度学习推理场景，而更大的任务量如训练中采用多个GPU，甚至多机多卡进行reduce则不在本文的讨论范围。本文先重点探讨使用最多的thread, warp, thread block reduce场景，其他后续慢慢更新。

不同协作方式的使用场景和原因

采用不同协作方式的原因主要有两点：

一是为了高效的内存访问，因为reduce类型算子一般是访存密集型而不是计算密集型，性能被访存而不是计算限制，如何高效的访问内存对于reduce算子至关重要。GPU的合并内存访问要求一个warp的线程访问的是内存连续地址的数据。如果不能保证一个warp保证访问的是连续的地址，至少也应该尽量保证sub-warp的4-8个线程访问的是连续的地址。

二是为了GPU线程更好的任务分配，充分利用硬件的并行性。例如避免创建过少的线程或者过多的线程，避免大量线程空转，避免单个线程循环次数过大。

题外话-基于向量数据类型的内存访问和计算

标量数据类型：例如float, int, half等等。

向量数据类型：float2, float4, half2, half4等等。cuda当前只提供了half和bf16的2元素向量数据类型，而half4等在OpenCL等非常常用，CUDA使用这种数据类型可以考虑参考half2自行扩展。

基于向量数据类型的内存访问和计算可以带来两个明显的好处：
1是内存访问时，只需要计算一次地址就可以访问得到多个数据，地址计算相应的指令占比明显降低。2是在循环任务中，向量化访问和计算可以成倍降低循环次数。例如一个warp对768长度的元素做reduce，那么标量访问需要循环24次，但是采用float4/half4向量访问和计算，那么只需要循环6次，那么内存访问和计算的指令占比明显提升，而循环本身的开销占比被明显降低。

reduce类型算子需要对每个线程循环处理多个数据，因此向量化访问和计算非常重要，但是到底是用4元素向量还是2元素向量，或者甚至更高的8元素向量，取决于任务场景和实际profiling性能。有时候并不见得half4一定性能比half2更高。

参考案例

https://github.com/NVIDIA/cuda-samples/tree/master/Samples/2_Concepts_and_Techniques/reduction

https://github.com/NVIDIA/cuda-samples/tree/master/Samples/2_Concepts_and_Techniques/reductionMultiBlockCG

https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cooperative-groups

https://developer.nvidia.com/blog/cooperative-groups/

thread reduce

基于一个线程进行reduce，通常针对两种场景，一个是对axis=-1 reduce但是reduce的元素个数只有少量几个的场景。另一个是非最内层维度reduce，例如[512,768]的张量沿着axis=0进行reduce，如下图。这种场景如果采用half2向量访问，那么可以分配768/2个线程，然后每个线程循环512次。这种方式保证了Coalesced memory accesses。

采用一个线程进行reduce的问题是可能存在循环次数过大的问题，如果reduce的元素很多，那么可以考虑多个thread block进行合作。

warp reduce

warp,sub-warp, thread block reduce一般都用于张量最内层维度axis=-1的reduce场景，从而保证Coalesced memory accesses。

warp内部线程reduce到一起一般基于Warp Shuffle方法，基于寄存器交互数据，效率很高。类似于OpenCL具有sub_group_reduce_add这样方便的函数，cuda也提供了类似的__reduce_add_sync方法，不过只支持int/uint类型以及compute capability >= 8.x。

https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-reduce-functions

https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-shuffle-functions

实现方法参考：

template <class T>
__device__ __forceinline__ T warpReduceSum(unsigned int mask, T mySum) {
  for (int offset = warpSize / 2; offset > 0; offset /= 2) {
    mySum += __shfl_down_sync(mask, mySum, offset);
  }
  return mySum;
}

#if __CUDA_ARCH__ >= 800
// Specialize warpReduceFunc for int inputs to use __reduce_add_sync intrinsic
// when on SM 8.0 or higher
template <>
__device__ __forceinline__ int warpReduceSum<int>(unsigned int mask,
                                                  int mySum) {
  mySum = __reduce_add_sync(mask, mySum);
  return mySum;
}
#endif

shuffle函数：

T __shfl_sync(unsigned mask, T var, int srcLane, int width=warpSize);
T __shfl_up_sync(unsigned mask, T var, unsigned int delta, int width=warpSize);
T __shfl_down_sync(unsigned mask, T var, unsigned int delta, int width=warpSize);
T __shfl_xor_sync(unsigned mask, T var, int laneMask, int width=warpSize);

T can be int, unsigned int, long, unsigned long, long long, unsigned long long, float or double. With the cuda_fp16.h header included, T can also be __half or __half2. Similarly, with the cuda_bf16.h header included, T can also be __nv_bfloat16 or __nv_bfloat162.

注意可以支持fp16和bf16，以及2元素打包的向量类型，但是性能不一定比float好。

mask用于指示有效参与reduce的线程，所有线程都有效直接设0xffffffff即可。

reduce通常把一个warp结果reduce到第一个线程，如果其他线程也想要则可以再用一次warp suffule进行一个broadcast操作即可。

value = __shfl_sync(0xffffffff, value, 0, 32);//把 laneId == 0线程的value变量的值广播给其他线程

基于warp reduce的layernorm算子

这里展示一个基于warp reduce的layernorm算子：

template <int pack_size, int channel>
__global__ layernorm_warp(const half4 *d_in, const half4 *d_scale, const half4 *d_bias, half4 *d_out, int batch,
                          float eps) {
  constexpr int WARP_SIZE = 32;
  int gid = blockIdx.x * blockDim.x + threadIdx.x;
  int cur_batch = gid / WARP_SIZE; // >>5
  int warp_lane = gid % WARP_SIZE; // & 0x1f
  if (cur_batch >= batch) {
    return;
  }
  constexpr int channel_pack = channel / pack_size;
  constexpr int loop = channel_pack / WARP_SIZE; // >>5

  int row_offs = cur_batch * channel_pack;

  half4 mean1_p = make_half4(CUDART_ZERO_FP16);
  half4 mean2_p = make_half4(CUDART_ZERO_FP16);

  for (int i = 0; i < loop; i++) {
    int addr = row_offs + i * WARP_SIZE + warp_lane;
    half4 x = d_in[addr];
    mean1_p += x;
    mean2_p += x * x;
  }
  float mean1_f = (mean1_p.x + mean1_p.y + mean1_p.z + mean1_p.w) / channel;
  float mean2_f = (mean2_p.x + mean2_p.y + mean2_p.z + mean2_p.w) / channel;
  mean1_f = warp_reduce_sum_broadcast(mean1_f);
  mean2_f = warp_reduce_sum_broadcast(mean2_f);
  float var = mean2_f - mean1_f * mean1_f;
  half std = sqrtf(var + eps);
  half mean = mean1_f;

  for (int i = 0; i < loop; i++) {
    int chn_id = i * WARP_SIZE + warp_lane;
    int addr = row_offs + chn_id;
    half4 x = d_in[addr];
    half4 tmp = (x - mean) / std;
    d_out[addr] = tmp * d_scale[chn_id] + d_bias[chn_id];
  }
}

当然这个kernel不一定是最优的，可以多方面进行改进：

1，这里用了half4访存和计算，但是half4不一定是最好的，实际可能half2还更好一些，但是直接half通常性能比较差。可以很容易修改权重模板定义支持各种向量类型。

2，采用warp那么总的线程数为batch*warp_size。但是一个thread block应该采用多少线程是最优的，256还是512等等，这是需要tune的。

3, 这里面部分转换成float，是为了部分提升准确度。但是为了极致的性能，可以考虑都half，但是溢出的风险更大，可能导致结果准确度不足。同时为了更好的准确度，元素和元素平方的累加也可以转换为float。

4，这里要读取2遍Input数据，可以考虑采用寄存器或者shared memory进行缓存。常见大模型的layernorm channel只有768,1024量级，放到一个warp 32个线程的寄存器还是足够的。当然即使不缓存，L1, L2 cache应该也能发挥其作用。

sub-warp reduce

sub-warp reduce是在warp reduce上的扩展，主要用于reduce的元素数量不是特别大的情况。

例如shape为[1024, 64]的张量最内层维度，使用half4向量访问时只需要16次循环，分配给一个warp 32个线程有一点浪费。而使用4-8个线程进行reduce即可。

实现sub-warp reduce只需要对warp reduce的warp size generalized到一个更小的size，把warpReduceSum稍微修改一下即可适配。

thread block reduce

一个thread block有多个warp，warp内部基于shuffle方法reduce，然后warp之间需要基于shared memory交换数据。

把一个thread block所有线程持有的元素reduce在一起一般分为3步：1，每个warp采用上面的warp reduce得到一个结果；2，每个warp 汇总结果的那个线程例如线程0把结果根据warp在thread block的id写入shared memory；3，然后采用一个warp例如第0个warp（当然也可以所有warp都读取也无所谓）读取shared memory，然后再做一次warp reduce。如果每个线程需要这个结果，那么再加一步broadcast的操作，例如如果步骤3只用warp0进行最终reduce则可以再把结果写入shared memory，然后所有线程读取这个shared memory。或者步骤3所有warp都读取shared memory进行reduce，那么只需要采用一个warp broadcast即可。

这里需要注意的一个点是一个thread block的线程数最好不超过1024，因为<=1024线程对应小于32个warp。如果超过1024线程，那么大于32个warp的结果在上面步骤3无法使用一个warp一次性进行reduce，至少要加一个循环之类的。

如果不确定一个block里面有多少个warp，可以定义动态大小的shared memory大小。然后通过kernel第三个lunch参数来指定：extern float __shared__ sdata[];

  extern float __shared__ sdata[];
  mySum = warpReduceSum<T>(mask, mySum);

  // each thread puts its local sum into shared memory
  if ((tid % warpSize) == 0) {
    sdata[tid / warpSize] = mySum;
  }

  __syncthreads();

  const unsigned int shmem_extent =
      (blockSize / warpSize) > 0 ? (blockSize / warpSize) : 1;
  const unsigned int ballot_result = __ballot_sync(mask, tid < shmem_extent);
  if (tid < shmem_extent) {
    mySum = sdata[tid];
    // Reduce final warp using shuffle or reduce_add if T==int & CUDA_ARCH ==
    // SM 8.0
    mySum = warpReduceSum<T>(ballot_result, mySum);
  }

  // write result for this block to global mem
  if (tid == 0) {
    g_odata[blockIdx.x] = mySum;
  }

多个线程块reduce

多个线程块之间进行reduce，首先仍然是对每个线程块采用thread block reduce，然后再对线程块之间的数据进行reduce。过去的GPU对于线程块之间传递数据只能基于global memory。这个额外的global memory哪里来，要么可以考虑临时把算子的output张量内存用来做一个临时的缓冲区来使用，要么就要单独给算子申请一个workspace内存。最新的nvidia Hopper架构增加了一个线程块集群的概念，也就是在grid和thread block之间额外引入了一个管理层级。在这个线程块集群里面的线程块之间可以基于分布式共享内存传递数据，相比global memory更快。

对于早于hopper架构的GPU来说，thread block之间reduce还需要采用协作组grid sync来进行同步，并且需要cudaLaunchCooperativeKernel来启动kernel。这个对于kernel能够创建的线程块数量有一定限制，应该是所有线程块能够同时驻留在GPU上。

参考样例：

cuda-samples-master\Samples\2_Concepts_and_Techniques\reductionMultiBlockCG

#include <stdio.h>
#include <iostream>
using namespace std;

#include <cuda_runtime.h>
#include "utils/cuda_mem_helper.h"
#include "utils/cuda_event_helper.h"
#include "utils/cuda_stream_helper.h"
#include "cuda_fp16.h"
#include "helper_math.h"

#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>

namespace cg = cooperative_groups;

template <class T>
__device__ __forceinline__ T warpReduceSum(unsigned int mask, T mySum) {
  const int WARP_SIZE = 32;
#pragma unroll
  for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
    mySum += __shfl_down_sync(mask, mySum, offset);
  }
  return mySum;
}

template <class T>
__device__ __forceinline__ T blockReduceSum(T* sdata, T data) {
  int tid = threadIdx.x;
  int warp_lane = tid % warpSize;
  int warp_id = tid / warpSize;
  int block_size = blockDim.x;
  int blk_warp_num = block_size / warpSize;

  data = warpReduceSum<T>(0xffffffff, data);

  // each thread puts its local sum into shared memory
  if (warp_lane == 0) {
    sdata[warp_id] = data;
  }
  __syncthreads();

  const unsigned int ballot_result = __ballot_sync(0xffffffff, tid < blk_warp_num);
  if (tid < blk_warp_num) {
    data = sdata[tid];
    data = warpReduceSum<T>(ballot_result, data);
  }
  return data;
}

const int co_block_num = 4;

__global__ void reduce_mblk(const float* __restrict__ d_in, float* __restrict__ d_out, int batch, int channel) {
  // used to reduce mean and variance togather
  __shared__ float sdata[32];

  int tid = threadIdx.x;
  int bid = blockIdx.x;
  int block_size = blockDim.x;

  int warp_lane = tid % warpSize;
  int warp_id = tid / warpSize;
  int blk_warp_num = block_size / warpSize;

  int gid = bid * block_size + tid;

  float data = d_in[gid];
  int grid_size = gridDim.x;

  data = blockReduceSum<float>(sdata, data);

  if (tid == 0) {
    printf("block_sum: %d %f\n", gid, data);
  }

  cg::grid_group grid = cg::this_grid();

  d_out[bid] = data;
  // sync of grid
  cg::sync(grid);

  float grid_sum = 0.0f;

  if (gid == 0) {
    for (int i = 0; i < grid_size; i++) {
      grid_sum += d_out[i];
    }
  }
  if (gid == 0) {
    printf("grid_sum: %d %f\n", gid, grid_sum);
  }
}

int main(void) {
  cudaError_t err = cudaSuccess;

  int batch = 1;
  int channel = 2048;

  CudaMemoryHelper<float> data_in({batch, channel});
  CudaMemoryHelper<float> data_out({batch, channel});

  // data_in.ReadFromFile("/mnt/e/quality_enhance/onnx/group_norm/test_data1/input.npy.bin");
  // data_mul.ReadFromFile("/mnt/e/quality_enhance/onnx/group_norm/test_data1/scale.npy.bin");
  // data_add.ReadFromFile("/mnt/e/quality_enhance/onnx/group_norm/test_data1/bias.npy.bin");
  data_in.StepInitHostMem(1.0f, 0.0f);
  data_in.CopyMemH2D();

  CudaStreamHelper stream_helper;
  CudaEventHelper event_helper(stream_helper.stream);

  int block_size = 512;
  int grid_size = batch * co_block_num;
  printf("CUDA kernel launch with %d blocks of %d threads\n", grid_size, block_size);

  int eval_num = 1;

  int smemSize = 0;
  void* kernelArgs[] = {(void*)&data_in.d_mem, (void*)&data_out.d_mem, (void*)&batch, (void*)&channel};

  dim3 dimBlock(block_size, 1, 1);
  dim3 dimGrid(grid_size, 1, 1);

  for (int i = 0; i < eval_num; i++) {
    // sync grid must use cudaLaunchCooperativeKernel to lunch kernel
    cudaLaunchCooperativeKernel((void*)reduce_mblk, dimGrid, dimBlock, kernelArgs, smemSize, stream_helper.stream);
    // reduce_mblk<<<grid_size, block_size, 0, stream_helper.stream>>>(data_in.d_mem, data_out.d_mem, batch, channel);
  }

  stream_helper.Sync();

  err = cudaGetLastError();
  if (err != cudaSuccess) {
    fprintf(stderr, "failed to launch cuda kernel: %s\n", cudaGetErrorString(err));
  }
  data_out.CopyMemD2H();
  // data_out.DumpToFile("/mnt/e/quality_enhance/onnx/group_norm/test_data1/output_actual.bin");

  return 0;
}