并行编程实战——CUDA编程的Warp Shuffle

一、线程间数据交互

在每一种语言中，对线程或进程间的数据交互都控制的非常谨慎。也就是说，为了兼顾效率和安全，往往对线程间的通信根据不同的情况提供不同的处理机制。特别涉及到内存中的数据交互，提供了多种的处理方法，典型的有：

共享内存
这种方式一般用来处理大块的数据交互，这种情况下，共享内存的优势还是比较大的
同步变量（含原子变量）
这种一般用来处理较小的数据，应用相对灵活
消息或事件
这种处理方式更加灵活并可以解耦多个模块间的数据交互，但其也只能提供中等规模的数据通信量

同样，在CUDA的多线程开发中，也会涉及到线程间的数据交互和通信，它也需要CUDA框架提供相关的数据交互的机制来满足实际开发的要求。

二、Warp Shuffle

在前面提到过CUDA中的共享内存等机制，此处分析一下Warp Shuffle，束内洗牌。它是一组强大的指令，提供了在同一个线程束（Warp）内的线程间直接互相访问彼此的寄存器的值的机制。简要的来说，就是在Warp内部的线程间进行变量交换。相对于共享内存而言，它是一种更加灵活、高效的多线程间数据交互的方法。

在这里不得不重新提到lane（通道）这个名词，它的意思是"A lane simply refers to a single thread within a warp"，其索引的Range为[0,31]。需要说明的是，同一个线程块可能有多个Warp，所以有可能会出现相同的lane index。

在CUDA中提供了多个Warp Shuffle Functions，不过在不同的版本可能有所不同。Warp Shuffle早期（计算力大于3.0）提供的接口在CUDA9.0(计算力大于7.0)后进行了更新，其提供了带_sync同步后缀的新接口。

Warp Shuffle主要特点有三个：首先是性能好，延迟低，寄存器的操作可以说是最快的行为之一了；其次通过使用寄存器替换共享内存间接的减少了内存的开销；最后，由于其内部提供了相关的同步语义，所以不需要开发者再显式的提供同步机制。

需要注意的是，Warp Shuffle只针对活动的线程且必须在同一个Warp中。如果操作的为非活动的线程或非同一个线程束内，则相关操作未定义。

三、相关的函数接口

CUDA主要提供了四种Warp Shuffle的函数，每种函数接口都对应一种特定的数据寻址模式（本文将不对早期的版本中的相关函数进行说明）。下面分别进行介绍：

T __shfl_sync(unsigned mask, T var, int srcLane, int width=warpSize)
在同一线程束内的分组指定下读取srcLane逻辑ID的变量Var的值。srcLane，获取数据的Lane（针对分组后的ID），注意width有默认值为线程束的大小。如果为下面的设定__shfl_sync(mask, var, 2)，默认width为线程束大小，则表示把当前线程束划分成一个逻辑组，它们都从线程2（即第3个线程）中的var中读取值。
T __shfl_up_sync(unsigned mask, T var, unsigned int delta, int width=warpSize)
在同一线程束内的分组指定下，每个线程向上（组内自身逻辑ID为基础）delta的ID线程获取var的值。delta，组内逻辑ID的距离值。在"自身ID+delta"范围外（组内，小于或大于），其线程的var保持不变。比如__shfl_up_sync(0xffffffff,x,2,16),表示所有线程都参与，划分为2个组，每组有16个线程。第1组从线程组内逻辑索引2（组内逻辑ID与线程束内线程ID相同，即第3个线程）获取x值；第2组从组内线程索引2（线程束内索引18，第19个线程）
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)
在同一线程束内的分组指定下，通过将当前线程的ID与laneMask进行XOR（异或）来计算源线程ID并获取其var值。laneMask，组内需要异或的线程逻辑ID。比如__shfl_xor_sync（0xffffffff,x,2,16）。它同样将线程束分成2组，每组16个线程。在每个组内与组内的线程ID 2 进行异或，获取其var的值。
需要说明的是，进行逻辑动作时，都是使用组内的逻辑ID，比如第2组实际的线程索引[16,31],但计算XOR时，仍然是使用[0,15],但在运算出来后，再加上基础的索引来求出真正的实际源线程ID，即15+XOR结果。
其实它的计算结果就是一个蝶形的分布，有兴趣的可以自己算一算或者上机跑一跑就明白了。

上面的接口可以发现它们三个相同的函数参数，下面进行统一说明：

mask
一个32位无符号整数，每一位对应着一个Warp中的32个线程。通过置0或1来表示线程的是否参与
var
这个好理解，就是要交换实际的变量的值，支持多种数据类型，如int,double等
width
Warp内的分组后组内线程的数量，必须是2的幂。Shuttle的操作只能在这个分出来的组内进行。比如是32，表示只有一个组，有32个线程用来交换数据。16表示分成两个组，每16个线程一组在组内进行交换。

在CUDA的应用中，上述的四个Shuttle函数，主要用于算法的处理。常见的如前面提到的Scan（扫描即前缀和）、相邻数据的交换以及广播和归约等。比如__shfl_xor_sync函数，在FFT和reduction中就非常常见。

四、例程

下面看官网提供的相关例程：

__shfl_sync

c 复制代码

#include <stdio.h>

__global__ void bcast(int arg) {
    int laneId = threadIdx.x & 0x1f;
    int value;
    if (laneId == 0)        // Note unused variable for
        value = arg;        // all threads except lane 0
    value = __shfl_sync(0xffffffff, value, 0);   // Synchronize all threads in warp, and get "value" from lane 0
    if (value != arg)
        printf("Thread %d failed.\n", threadIdx.x);
}

int main() {
    bcast<<< 1, 32 >>>(1234);
    cudaDeviceSynchronize();

    return 0;
}

__shfl_up_sync

c 复制代码

#include <stdio.h>

__global__ void scan4() {
    int laneId = threadIdx.x & 0x1f;
    // Seed sample starting value (inverse of lane ID)
    int value = 31 - laneId;

    // Loop to accumulate scan within my partition.
    // Scan requires log2(n) == 3 steps for 8 threads
    // It works by an accumulated sum up the warp
    // by 1, 2, 4, 8 etc. steps.
    for (int i=1; i<=4; i*=2) {
        // We do the __shfl_sync unconditionally so that we
        // can read even from threads which won't do a
        // sum, and then conditionally assign the result.
        int n = __shfl_up_sync(0xffffffff, value, i, 8);
        if ((laneId & 7) >= i)
            value += n;
    }

    printf("Thread %d final value = %d\n", threadIdx.x, value);
}

int main() {
    scan4<<< 1, 32 >>>();
    cudaDeviceSynchronize();

    return 0;
}

__shfl_xor_sync

c 复制代码

#include <stdio.h>

__global__ void warpReduce() {
    int laneId = threadIdx.x & 0x1f;
    // Seed starting value as inverse lane ID
    int value = 31 - laneId;

    // Use XOR mode to perform butterfly reduction
    for (int i=16; i>=1; i/=2)
        value += __shfl_xor_sync(0xffffffff, value, i, 32);

    // "value" now contains the sum across all threads
    printf("Thread %d final value = %d\n", threadIdx.x, value);
}

int main() {
    warpReduce<<< 1, 32 >>>();
    cudaDeviceSynchronize();

    return 0;
}

给出最后一个的运行结果：

复制代码

Thread 0 final value = 496
Thread 1 final value = 496
Thread 2 final value = 496
Thread 3 final value = 496
Thread 4 final value = 496
Thread 5 final value = 496
Thread 6 final value = 496
Thread 7 final value = 496
Thread 8 final value = 496
Thread 9 final value = 496
Thread 10 final value = 496
Thread 11 final value = 496
Thread 12 final value = 496
Thread 13 final value = 496
Thread 14 final value = 496
Thread 15 final value = 496
Thread 16 final value = 496
Thread 17 final value = 496
Thread 18 final value = 496
Thread 19 final value = 496
Thread 20 final value = 496
Thread 21 final value = 496
Thread 22 final value = 496
Thread 23 final value = 496
Thread 24 final value = 496
Thread 25 final value = 496
Thread 26 final value = 496
Thread 27 final value = 496
Thread 28 final value = 496
Thread 29 final value = 496
Thread 30 final value = 496
Thread 31 final value = 496

五、总结

warp shuffle更像是一种基础算法的固件支持接口。这些基础的算法提供了强大的算力供上层的算法使用。通过基础的算法单元快速进行多种算法的实现，达到了灵活而又高效的设计目的。