warp shuffle指令

目录

abstract

背景

线程束洗牌指令的不同形式

线程束内的共享内存数据

线程束值的广播

线程束上偏移

线程束下偏移

线程束内环绕移动

线程束内的蝴蝶交换

线程束内的归约

总结

---

abstract

本文将介绍线程束洗牌指令

前文我们介绍了共享内存，常量内存，只读缓存，今天将介绍一个特殊的机制-warp shuffle

背景

之前不断提到过寄存器是每个线程独有的，也就是不能窥见对方，我们的所有数据都必须加载到寄存器里面，cuda核心才能计算

warp shuffle提供的能力：可以相互访问对方的寄存器，这样提供了一种能够交换数据的能力

这样交换数据不必通过l1缓存或者共享内存，效率更高

总结：线程束洗牌指令是线程束内线程通信的极佳方式

线程束洗牌指令的不同形式

这里需要计算线程属于哪个线程束，我们之前提到过threadIdx，这个是block内的idx

所以判断这个线程在block的哪个线程束，在线程束当中的哪个位置

int LaneId=threadIdx.x%32;    //属于线程束当中哪个线程
int WarpId=threadIdx.x/32;    //属于哪个线程束

// 1 __shfl_sync 按源线程id读取值（广播最常用）
T __shfl_sync(unsigned mask, T val, int src_lane, int width=32);

// 2 __shfl_up_sync 向上移位：lane = lane - delta
T __shfl_up_sync(unsigned mask, T val, int delta, int width=32);

//3 __shfl_down_sync 向下移位：lane = lane + delta
T __shfl_down_sync(unsigned mask, T val, int delta, int width=32);

//4 __shfl_xor_sync 按lane id做异或，蝴蝶模式，归约经典
T __shfl_xor_sync(unsigned mask, T val, int lane_mask, int width=32);

这些函数后面都跟__sync是因为新的架构有独立线程调度机制，所以可能会出现不同步的情况，这些函数就是强制同步，然后再往下执行，老的架构是天然同步的

• mask的意思就是活跃掩码，一个32位，第几位为1就表示第几个lane线程参与，比如0xffffffff就表示一个warp内的32线程都要参与这条指令，也就是线程束洗牌指令，强制要求这些参加的thread到达一个shuffle点，这个是同步点，然后才能往下交换
• val就是当前线程要贡献出去的值，就是把寄存器的值贡献出去，shuffle的作用就是把某个线程的val交换给别人
• src_lane：表示要将哪个线程的val贡献出去，这就是定位哪个源线程的
• width：虚拟 Warp 大小，默认值 32。它把整个 Warp 分成若干个大小都为 width 的组 ，srcLane 是在这个组内的相对索引。（比如你的参数是16，那就是分为2组，然后每组里面的src_lane就是源线程，假设是2，那就是2和18提供val给各自的组，组之间不可以跨组交换）
• delta：向上偏移量，目标线程=当前线程的lane-delta（当前 lane = 5，delta = 2，width = 16。目标线程 = 5 - 2 = 3。所以 lane 5 的线程会得到 lane 3 的 var 值。如果目标线程编号小于 0（即 lane - delta < 0），则该线程保持自己的原值不变。）
• 那向下偏移量：目标=当前lane+delta（如果目标线程编号超出当前组（lane + delta >= width），则该线程保持原值不变）
• lane_mask：异或掩码。目标线程的 lane 编号 = 当前线程的 lane 编号 异或（XOR） laneMask。laneMask = 1，二进制是 000001。当前 lane 和 1 异或后的结果就是目标 lane。这样效果是：lane 0 和 lane 1 互相交换，lane 2 和 lane 3 互相交换......这在归约和蝴蝶交换算法中非常常用。
• 返回值就是val这个变量的值

由于以上的所有参数都是模板，所以可以是int float等等，val是一个变量名哦，我们提前设了变量，比如每个线程的 int val=threadIdx.x，那交换就是把目标线程的val这个变量里面的值拿出来给大家交换

这张图是谭升大佬画的，所以很详细，是关于__shlf_xor_sync的lane_mask=1，也就是相邻两两互相交换

线程束内的共享内存数据

接下来我们用代码实现以下，看看每一个指令的作用效果，洗牌指令可以用于下面三种变量类型中：

• 标量变量
• 数组
• 向量型变量

向量型：一次打包，减少指令条数

类型	包含的标量	总字节数
int2 / float2	2 个 int / float	8 字节
int3 / float3	3 个 int / float	12 字节
int4 / float4	4 个 int / float	16 字节

线程束值的广播

#include <iostream>
#include "../freshman.hpp"
using namespace std;
#define BDIM 16
#define SCRLANE 2
__global__ void shuffle_broadcast(int* out,int* in,int const scrLane){
    int var=in[threadIdx.x];
    var=__shfl_sync(0xffffffff,var,scrLane,BDIM); //每个线程都将获取到scrLane的var
    out[threadIdx.x]=var;                       //改变传出去
}
int main(){
    int device=0;
    CHECK(cudaSetDevice(device));
    int Elem=1<<4;
    int* in_h;
    int* out_dtoh;
    in_h=(int*)malloc(sizeof(int)*Elem);
    out_dtoh=(int*)malloc(sizeof(int)*Elem);
    for(int i=0;i<Elem;i++){
        in_h[i]=i;
    }
    int* in_d=nullptr;
    int* out_d=nullptr;
    CHECK(cudaMalloc((void**)&in_d,sizeof(int)*Elem));
    CHECK(cudaMalloc((void**)&out_d,sizeof(int)*Elem));
    CHECK(cudaMemcpy(in_d,in_h,sizeof(int)*Elem,cudaMemcpyHostToDevice));

    dim3 block(BDIM);
    dim3 grid((Elem+block.x-1)/block.x);
    
    shuffle_broadcast<<<grid,block>>>(out_d,in_d,SCRLANE);
    cudaDeviceSynchronize();
    cudaMemcpy(out_dtoh,out_d,sizeof(int)*Elem,cudaMemcpyDeviceToHost);
    for(int i=0;i<Elem;i++){
        cout<<out_dtoh[i]<<" ";
    }
    cout<<endl;
    free(in_h);
    free(out_dtoh);
    cudaFree(in_d);
    cudaFree(out_d);
    return 0;
}

这个就是最简单的，就是把线程束内的2号线程的var变量的值换过来

线程束上偏移

__global__ void shuffle_up_broadcast(int* out,int* in,int const scrLane){
    int var=in[threadIdx.x];
    var=__shfl_up_sync(0xffffffff,var,scrLane,BDIM); //每个线程都将获取到scrLane的var
    out[threadIdx.x]=var;                       //改变传出去
}

那头两个线程是会越界的，所以他们的值不变，所以答案就是0 1 0 1 2 3 4 5......

线程束下偏移

__global__ void shuffle_down_broadcast(int* out,int* in,int const scrLane){
    int var=in[threadIdx.x];
    var=__shfl_down_sync(0xffffffff,var,scrLane,BDIM); //每个线程都将获取到scrLane的var
    out[threadIdx.x]=var;                       //改变传出去
}

线程束内环绕移动

__global__ void shuffle_cycle_broadcast(int* out,int* in,int const offset){
    int var=in[threadIdx.x];
    var=__shfl_sync(0xffffffff,var,threadIdx.x+offset,BDIM); //每个线程都将获取到scrLane的var
    out[threadIdx.x]=var;                       //改变传出去
}

当我们的offset为2时，那就是每个都是获得本身往后移动两个的线程的元素，但是最后两个为什么是这样？原因在于取余了

目标lane = 当前组的起始lane + (srcLane % width)

其他就不取余了，要注意哦

线程束内的蝴蝶交换

废话不多说，上代码

__global__ void shuffle_xor_broadcast(int* out,int* in,int const offset){
    int var=in[threadIdx.x];
    var=__shfl_xor_sync(0xffffffff,var,1,BDIM); //每个线程都将获取到scrLane的var
    out[threadIdx.x]=var;                       //改变传出去
}

线程束内的归约

__inline__ __device__ int warpReduce(int localSum)
{
    localSum += __shfl_xor(localSum, 16);
    localSum += __shfl_xor(localSum, 8);
    localSum += __shfl_xor(localSum, 4);
    localSum += __shfl_xor(localSum, 2);
    localSum += __shfl_xor(localSum, 1);
    return localSum;
}
__global__ void reduceShfl(int * g_idata,int * g_odata,unsigned int n)
{
	//set thread ID
    __shared__ int smem[DIM];
	unsigned int idx = blockDim.x*blockIdx.x+threadIdx.x;
	//convert global data pointer to the
    //1.
	int mySum=g_idata[idx];
	int laneIdx=threadIdx.x%warpSize;
	int warpIdx=threadIdx.x/warpSize;
    //2.
	mySum=warpReduce(mySum);
    //3.
	if(laneIdx==0)
		smem[warpIdx]=mySum;
	__syncthreads();
    //4.
	mySum=(threadIdx.x<DIM)?smem[laneIdx]:0;
	if(warpIdx==0)
		mySum=warpReduce(mySum);
    //5.
    if(threadIdx.x==0)
		g_odata[blockIdx.x]=mySum;

}

第一轮我们让每个线程束内的线程都去找别的线程相加，得到的结果是warp内每个线程的localSum都一样

然后我们取第一个即可，放到共享内存当中，因为warp之间要交换数据，所以我们让每个warp的第一个线程把localSum的数据写到共享内存即可

再使用第一个warp前面几个线程把共享内存的数据读到一个warp里面，这里面再进行一轮归约，注意只有前DIM个有数据，其余的为0

也就是第一个warp内的数据为：sum sum sum sum sum sum sum sum sum 0 0 0 0 0......

这样0不影响，只要把这些再归约一遍，然后输出即可

可以实验一下，发现warp shuffle的方法是最快的，比共享内存还快，更不用说全局了

总结

这一章的核心思想是：Warp Shuffle 是当前 GPU 上最快的线程间通信方式，因为它直接在寄存器之间交换数据，完全跳过了共享内存和全局内存。

Warp Shuffle 就是目前归约优化的"终极形态"，它通过跳过所有中间内存，直接在寄存器间交换数据，达到了极致的低延迟和高带宽。这也是后续学习高性能库（如 CUB、cuBLAS）中许多优化技巧的基石。