最近写算子发现 reduce 和 prefix 还是有点像的,都需要过一遍数据，reduce 把结果汇集到一个结果，prefix则在过程中维护每个位置的结果

它们都是一个类似树的过程。

reduce sum 过程

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
+----+----+----+----+----+----+----+----+
| a  | b  | c  | d  | e  | f  | g  | h  |
+----+----+----+----+----+----+----+----+
   \   /     \   /     \   /     \   /
    \ /       \ /       \ /       \ /
     |         |         |         |
     v         v         v         v
+---------+---------+---------+---------+
|  a + b  |  c + d  |  e + f  |  g + h  |
+---------+---------+---------+---------+
        \   /               \   /
         \ /                 \ / 
          |                   |  
          v                   v  
     +---------+       +---------+
     | a+b+c+d |       | e+f+g+h |
     +---------+       +---------+
                 \   /
                  \ / 
                   |  
                   v  
        +-----------------+
        | a+b+c+d+e+f+g+h |
        +-----------------+

这么看是个树，每次加的是相邻的2个，也可以每次隔stride相加，即[(1,4),(2,5),(3,6),(4,7)],[(1,3),(2,4)],[(1,2)]

1
2
3
4
5
6
7
float v = arr[idx];
unsigned int mask = 0xffffffff;
for(int stride = 16; stride > 0; stride >>= 1){
     v += __shfl_down_sync(mask, v, stride);
     mask >>= 1;
}
arr[idx] = v;

prefix sum 过程

每个元素i加到i+stride位置，stride = 1,2,4,8…2^n
得到结果
CDUA warp shuffle 很好实现
prefix sum

1
2
3
4
5
6
7
8
float v = arr[idx];
for(int stride = 1; i < 32; stride >>= 1){
     float t = __shfl_up_sync(0xffffffff, v, stride);
     if(idx & 31 >= stride){
          v += t;
     }
}
arr[idx] = v;

idx是当前全局数组的index
stride是位移量
32 是Warp大小
idx & 31 是 lane_id，warp内每个线程为lane

可以改一下直接用mask过滤

1
2
3
4
5
6
7
float v = arr[idx];
unsigned int mask = 0xffffffff;
for(int stride = 1; i < 32; stride >>= 1){
     mask >>= 1;
     v += __shfl_up_sync(mask, v, stride);
}
arr[idx] = v;