ROCm Reduce算子（Warp Shuffle/SMEM规约）

Reduce

规约操作，实际上就是把一个张量的所有元素，汇总到一个值。汇总可以是求和，求max。这里以求和为例。

本文研究一下ROCm版Reduce的两种写法，以及分析线程束分化对性能的影响。

写法1：块内同步

每个线程块负责一个区域，整个线程块内的所有线程一起写作，由于执行基本单位是线程束不是线程块，块内的不同线程束之间是异步的，需要加上强制同步屏障。

• extern __shared__ float sdata[];动态共享内存，编译时不确定，通过kernel launch时的参数，运行时确定。
• int global = blockIdx.x * blockDim.x * 2 + tid;计算当前线程的全局编号
• if (global < n) sum += in[global];用全局编号，去全局内存中的初始数据读取
• if (global + blockDim.x < n) sum += in[global + blockDim.x];经典跨步读取，一个线程会读取多个位置，提高计算密度。这里相当于每个线程块，会负责blockDim.x * 2长度的一个区间。
• sdata[tid] = sum;读到的数据存入共享内存，等待后续规约
• __syncthreads();块内所有线程都读完才能继续，但是块内不同线程束是不同步的，需要强制同步
• for (int stride = blockDim.x / 2; stride > 0; stride >>= 1)经典规约算法，类似于log x次操作求出x位整形的1个数奇偶，每次都把有效范围内后一半元素折叠到前一半。比如64个元素的话，开始把[32,63]折叠到[0,31]，有效范围变为[0,31]，然后下一步把[16,31]折叠到[0,15]，以此类推直到有效元素只剩一个。
• if (tid < stride)为了保证每次都折叠到前一半，需要if else判断线程是否位于前一半，这就引入了线程束分化
• __syncthreads();同样吗，每一轮折叠必须所有线程都完成了才能进入下一轮，但不同线程束是异步的，需要强制同步
• if (tid == 0) out[blockIdx.x] = sdata[0];我们最后把每个线程块的元素规约成了一个值，根据线程块编号写入答案数组。这里并没有一步到位把整个数组规约成一个元素，而是类似于一个递归，每次都把问题规模缩小，然后多次调用这个内核实现规约到只剩一个元素

这个内核比较简单，但是有明显的性能问题，一是线程束分化if (tid < stride)，造成分化的线程束并行度降低，退化成串行。而是频繁的__syncthreads();使得很多已经完成工作的线程处于阻塞，必须等待少数几个未完成的线程。如何避免这两点？我们来看下一个内核。

// 版本1：分支+同步的传统 reduction
__global__ void reduce_branchy(const float* __restrict__ in,
float* __restrict__ out,
int n) {
extern __shared__ float sdata[];
int tid = threadIdx.x;
int global = blockIdx.x * blockDim.x * 2 + tid;

float sum = 0.0f;
if (global < n) sum += in[global];
if (global + blockDim.x < n) sum += in[global + blockDim.x];
sdata[tid] = sum;
__syncthreads();

for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
sdata[tid] += sdata[tid + stride];
}
__syncthreads();
}

if (tid == 0) out[blockIdx.x] = sdata[0];
}

写法2：Warp级规约

每个block不同步，但是每个线程束是同步的。那么能不能考虑把规约的基本单位设为一个线程束。每个线程束各自规约，把结果写到一个临时数组，最后再启动一个线程束把临时数组规约，得到整个块的结果

• wf_reduce_sum先实现一个线程束规约的函数
• for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1)内部折叠思路和线程块折叠类似，只是把范围改成从线程束大小开始，而不是块大小开始
• __shfl_down(v, offset, WARP_SIZE);规约操作和整块规约完全不同。同一线程束，有内置函数，不用访问共享内存。这个函数的意思是读取当前线程编号增加offset后的另一个线程的v变量，如果增加offset后超出warp_size则结果为0。这是利用了同一warp的线程之间，存在直接的数据通路，不用借助共享内存进行通信。
• int lane = tid % WARP_SIZE;当前线程在线程束内的编号。我们默认全局的所有线程，就是没warp_sizeg个划分成一个warp
• int wid = tid / WARP_SIZE;当前线程所在的线程束编号。
• float wf_sum = wf_reduce_sum(v);每个线程都一起执行线程束规约函数，由于线程束内是同步的，不需要同步屏障了
• if (lane == 0) wf_sums[wid] = wf_sum;线程束规约结束了，线程束内第一个线程负责把结果写入临时数组
• __syncthreads();这里由于不同线程束之间还是不同步的，需要一个同步屏障
• if (wid == 0)每个线程束的结果都写入数组了，让第一个线程束来对临时数组再次进行线程束规约，得到整个线程块的答案
• int num_waves = (blockDim.x + WARP_SIZE – 1) / WARP_SIZE;首先计算当前线程块内线程束个数
• float x = (lane < num_waves) ? wf_sums[lane] : 0.0f;如果线程束个数小于线程束内线程个数，多出来的线程也要参与计算，只是把初始数据赋值为0。由于是线程束内规约，不读SMEM，直接把数据放到一个寄存器变量里
• 再次执行线程束规约得到整块结果float block_sum = wf_reduce_sum(x);
• if (lane == 0) out[blockIdx.x] = block_sum;当前线程数的第一个线程写入全局答案数组

这个版本有几个有点：

• 一是整个warp内都是同步的，不需要额外的同步操作，
• 二是warp内操作都相同，没有线程束发散，
• 三是不用走SMEM，数据搬运更快。

// wavefront shuffle sum
__device__ __forceinline__ float wf_reduce_sum(float v) {
for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
v += __shfl_down(v, offset, WARP_SIZE);
}
return v;
}

__global__ void reduce_shuffle(const float* __restrict__ in,
float* __restrict__ out,
int n) {
extern __shared__ float wf_sums[];

int tid = threadIdx.x;
int idx = blockIdx.x * blockDim.x * 2 + tid;

float v = 0.0f;
if (idx < n) v += in[idx];
if (idx + blockDim.x < n) v += in[idx + blockDim.x];

float wf_sum = wf_reduce_sum(v);

int lane = tid % WARP_SIZE;
int wid = tid / WARP_SIZE;

if (lane == 0) wf_sums[wid] = wf_sum;
__syncthreads();

if (wid == 0) {
int num_waves = (blockDim.x + WARP_SIZE – 1) / WARP_SIZE;
float x = (lane < num_waves) ? wf_sums[lane] : 0.0f;
float block_sum = wf_reduce_sum(x);
if (lane == 0) out[blockIdx.x] = block_sum;
}
}

多次调用

前面说过，这内核只是把每个block都规约成一个数，每做一次数组长度/block_size，想要实现规约到一个元素，需要对内核进行封装，循环调用，代码如下

• int max_blocks = (n + (threads * 2 – 1)) / (threads * 2);首先计算第一轮的block数，用来申请缓冲区，每个线程负责两个元素，所以块长是block线程数的两倍
• float* d_buf1 = nullptr;float* d_buf2 = nullptr;定义两个缓冲区，因为reduce的输入输出必须是两个数组，多次调用reduce至少需要两个缓冲区，来回倒着用。
• auto launch_once = [&](int cur_n, const float* cur_in, float* cur_out)一个lambda表达式，传入输入，输出缓冲区，以及当前的有效元素个数，执行一轮规约
• int blocks = (cur_n + (threads * 2 – 1)) / (threads * 2);计算当前的block数
• if (use_shuffle)判断使用哪个内核
• smem = num_waves * sizeof(float);根据每个block启动的warp数，确定每个block所需的共享内存大小，前面我们声明的extern SMEM会被初始化为这个大小
• int next_n = launch_once(cur_n, cur_in, cur_out);启动一次的函数，返回值是这次规约后的有效元素个数，将这个作为下一轮的规约长度
• cur_in = cur_out;这一轮的规约结果，作为下一轮的规约输入
• cur_out = (cur_out == d_buf1) ? d_buf2 : d_buf1;下一轮使用和这一轮不一样的数组作为输出缓冲区，类似于滚动数组优化的思路
• while (cur_n > 1) 重复规约直到有效元素只剩一个

// Host：多轮 reduction 直到剩一个数
float run_reduce(const float* d_in, int n,
bool use_shuffle,
int threads, int iterations) {
hipEvent_t start, stop;
HIP_CHECK(hipEventCreate(&start));
HIP_CHECK(hipEventCreate(&stop));

int max_blocks = (n + (threads * 2 – 1)) / (threads * 2);
float* d_buf1 = nullptr;
float* d_buf2 = nullptr;
HIP_CHECK(hipMalloc(&d_buf1, max_blocks * sizeof(float)));
HIP_CHECK(hipMalloc(&d_buf2, max_blocks * sizeof(float)));

auto launch_once = [&](int cur_n, const float* cur_in, float* cur_out) {
int blocks = (cur_n + (threads * 2 – 1)) / (threads * 2);
size_t smem = 0;
if (use_shuffle) {
int num_waves = (threads + WARP_SIZE – 1) / WARP_SIZE;
smem = num_waves * sizeof(float);
hipLaunchKernelGGL(reduce_shuffle, dim3(blocks), dim3(threads),
smem, 0, cur_in, cur_out, cur_n);
} else {
smem = threads * sizeof(float);
hipLaunchKernelGGL(reduce_branchy, dim3(blocks), dim3(threads),
smem, 0, cur_in, cur_out, cur_n);
}
return blocks;
};

// warmup
{
int cur_n = n;
const float* cur_in = d_in;
float* cur_out = d_buf1;
while (cur_n > 1) {
int next_n = launch_once(cur_n, cur_in, cur_out);
cur_n = next_n;
cur_in = cur_out;
cur_out = (cur_out == d_buf1) ? d_buf2 : d_buf1;
}
HIP_CHECK(hipDeviceSynchronize());
}

HIP_CHECK(hipEventRecord(start));
for (int it = 0; it < iterations; ++it) {
int cur_n = n;
const float* cur_in = d_in;
float* cur_out = d_buf1;

while (cur_n > 1) {
int next_n = launch_once(cur_n, cur_in, cur_out);
cur_n = next_n;
cur_in = cur_out;
cur_out = (cur_out == d_buf1) ? d_buf2 : d_buf1;
}
}
HIP_CHECK(hipEventRecord(stop));
HIP_CHECK(hipEventSynchronize(stop));

float ms = 0.0f;
HIP_CHECK(hipEventElapsedTime(&ms, start, stop));

HIP_CHECK(hipFree(d_buf1));
HIP_CHECK(hipFree(d_buf2));
HIP_CHECK(hipEventDestroy(start));
HIP_CHECK(hipEventDestroy(stop));
return ms;
}

性能分析