CUDA WarpReduce 学习笔记

前言之前看我司的 如何实现一个高效的Softmax CUDA kernel？多少还是有些细节没有理解，恰好最近要做一个类似的 Reduce+Scale Kernel，原理机制还是比较相似的，所以翻出来重新理解一下。背景我们定义这么一个ReduceScale操作：假设Tensor是(N, C)，首先在C这个维度计算出 absMax 值，我们记作scale，然后将每一行除以各自 行的scale，并最终输出。一段朴素的numpy代码是这样：import numpy as npN = 1000C = 128x = np.random.randn(N, C)scale = np.expand_dims(np.max(np.abs(x), axis=1), 1)out = x / scaleprint(out.shape)BaseLine这里我们BaseLine是直接调用cub库中的 BlockReduce，一个 threadBlock 处理一行数据，计算出AbsMaxVal，然后再缩放，代码如下：#include "cuda.h"#include "cub/cub.cuh"constexpr int kReduceBlockSize = 128;templatetypename T>__device__ T abs_func(const T& a) {  return abs(a);}templatetypename T>_ ………………………………