CUDA中Warp Shuffle的使用详解

创作时间:

2025-01-22 03:19:35

作者:

@小白创作中心

CUDA中Warp Shuffle的使用详解

在CUDA编程中，Warp Shuffle是一种非常重要的并行计算技术，它允许线程之间直接交换数据，而无需通过共享内存。本文将详细介绍Warp Shuffle的两个主要函数：__shfl_xor_sync和__shfl_sync，并通过代码示例展示它们的使用方法。

一 __shfl_xor_sync

template<typename T, const int kWarpSize = WARP_SIZE>
DEVICE_INLINE T warp_reduce_max(T val) {
  #pragma unroll
  for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {
    val = max(val, __shfl_xor_sync(0xffffffff, val, mask, kWarpSize));
  }
  return val;
}

__shfl_xor_sync() 通过对调用者的通道 ID 与 laneMask 执行按位异或来计算源通道 ID：返回结果通道 ID 所持有的 var 的值。如果宽度小于warpSize，那么每组宽度连续的线程都能够访问早期线程组中的元素，但是如果它们尝试访问后面线程组中的元素，则将返回他们自己的var值。这种模式实现了一种蝶式寻址模式，例如用于树规约和广播。

1.1 使用实例

1.1.1 理论解释

如果mask = 2，则：

图1.1：当mask=2时的线程数据交换示意图

说明：对于lane0，mask = 2 ，则逐位异或 0 ^ 2 = 2,1 ^ 2 = 3，和图2.2相符合
width把warp切成好几个group，我可以访问在我前面的group，但访问后面的group就会返回我自己的值

1.2.2 测试

#include <stdio.h>
__global__ void warpReduce()
{
    int laneId = threadIdx.x & 0x1f;
    // Seed starting value as inverse lane ID
    int value = laneId;
    // Use XOR mode to perform butterfly reduction
    for (int i = 2; 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;
}

asm代码为：

template<typename T, const int kWarpSize = WARP_SIZE>
DEVICE_INLINE T warp_reduce_sum(T val) {
  #pragma unroll
  for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {
    val += __shfl_xor_sync(0xffffffff, val, mask, kWarpSize);
  }
  return val;
}

reduce过程为：

图1.3：reduce过程示意图

本质上是4个一组求和（kWrapSize = 4和kWrapSize = 32）得到的结果是相同的

输出结果为

二 __shfl_sync

图2.1：__shfl_sync函数示意图

__shfl_sync() 返回由 srcLane 给定 ID 的线程持有的 var 的值。如果 width 小于 warpSize，则 warp 的每个子部分都表现为一个单独的实体，其起始逻辑通道 ID 为 0。如果 srcLane 超出范围 [0:width-1]，则返回的值对应于通过 srcLane srcLane modulo width所持有的 var 的值（即在同一部分内）。

2.1 实例

#include <stdio.h>
__global__ void bcast(int arg)
{
    int laneId = threadIdx.x & 0x1f;
    int value;
    // Note unused variable for all threads except lane 0
    // if (laneId == 0)
        // value = arg;
    value = laneId;
    __syncthreads();
    printf("before braodcast Thread %d value = %d\n", threadIdx.x,value);
    // Synchronize all threads in warp, and get "value" from lane 0
    value = __shfl_sync(0xffffffff, value, laneId +1 ,4);
    // if (laneId == 0)
    printf("after broadcast Thread %d value = %d\n", threadIdx.x,value);
}
int main()
{
    bcast<<<1, 32>>>(1234);
    cudaDeviceSynchronize();
    return 0;
}

对应的输出为: