CUDA 知识点：线程束洗牌函数 | CUDA

CUDA 中的线程束内基本函数包括：

• 线程束表决函数（warp vote functions）
• 线程束匹配函数（warp match functions）
• 线程束洗牌函数（warp shuffle functions）
• 线程束矩阵函数（warp matrix functions）

其中，线程束匹配函数和线程束矩阵函数都只能在 Volta 及更高架构的 GPU 中使用。本文主要介绍线程束洗牌函数。

基本概念

首先介绍 warp 中的一个概念：lane。一个 lane 就是一个 warp 中的一个 thread，每个 lane 在同一个 warp 中由 lane 索引唯一确定，因此其范围为 [0,31]。在一个一维的 block 中，可以通过下面的方式计算索引：

int lane_id = threadIdx.x % 32;
int warp_id = threadIdx.x / 32;

例如，在同一个 block 中的 thread_0 和 thread_32 的 lane_id 都为 0。

函数功能介绍

线程束洗牌函数的接口如下：

每个线程束洗牌函数的最后一个参数 w 都是可选的，有默认值 warpSize，它在当前所有架构的 GPU 中都是 32。参数 w 只能取 2、4、8、16、32 这 5 个整数中的一个。当 w 小于 32 时，就相当于（逻辑上的）线程束大小是 w，而不是 32，其他规则不变。这样的话上述计算索引的方式就要变为：

int lane_id = threadIdx.x % w;
int warp_id = threadIdx.x / w;

掩码用于指定将要参与计算的线程，当掩码中的一个二进制位为 1 时，代表对应的线程参与计算；当掩码中的一个二进制位为 0 时，代表忽略对应的线程。各种函数返回的结果对被掩码排除的线程来说是未定义的。

__shfl_sync

参与该次操作的所有线程通过该函数获得 lane_id 为 srcLane 的线程中变量 v 的值。这是一种广播式数据交换，即将一个线程中的数据广播到所有（包括自己）线程。

__shfl_up_sync

参与该次操作的所有线程将自己的 lane_id 减去 d 来计算源线程的 lane_id，这样源线程的 v 就会返回调用线程。warp 中最开始的 d 个的线程不会改变。

__shfl_down_sync

参与该次操作的所有线程将自己的 lane_id 加上 d 来计算源线程的 lane_id，这样源线程的 v 就会返回调用线程。warp 中最末尾的 d 个的线程不会改变。

__shfl_xor_sync

lane_id 为 t 的参与线程返回 lane_id 为 t^1aneMask 的线程中变量 v 的值。其中 ^ 表示两个整数按位做异或运算。该函数可以让线程束内线程两两交换数据。

参考

《CUDA 编程：基础与实践》（樊哲勇，清华大学出版社）

CUDA 之 Warp Shuffle 详解