如何避免这个 CUDA 内核中的线程分歧？

Question

对于CUDA核函数，得到如下所示的分支散度，如何优化？

int gx = threadIdx.x + blockDim.x * blockIdx.x;
val = g_data[gx];

if (gx % 4 == 0)
    val = op1(val);
else if (gx % 4 == 1)
    val = op2(val);
else if (gx % 4 == 2)
    val = op3(val);
else if (gx % 4 == 3)
    val = op4(val);

g_data[gx] = val;

Answer 1

如果我在 CUDA 中编程，我当然不会做这些。但是要回答您的问题：

如何避免这个 CUDA 内核中的线程分歧？

你可以这样做：

int gx = threadIdx.x + blockDim.x * blockIdx.x;
val = g_data[gx];

int gx_bit_0 =  gx & 1;
int gx_bit_1 = (gx & 2) >> 1;
val = (1-gx_bit_1)*(1-gx_bit_0)*op1(val) + (1-gx_bit_1)*(gx_bit_0)*op2(val) + (gx_bit_1)*(1-gx_bit_0)*op3(val) + (gx_bit_1)*(gx_bit_0)*op4(val);

g_data[gx] = val;

这是一个完整的测试用例：

$ cat t1914.cu
#include <iostream>

__device__ float op1(float val) { return  val + 1.0f;}
__device__ float op2(float val) { return  val + 2.0f;}
__device__ float op3(float val) { return  val + 3.0f;}
__device__ float op4(float val) { return  val + 4.0f;}

__global__ void k(float *g_data){

  int gx = threadIdx.x + blockDim.x * blockIdx.x;
  float val = g_data[gx];

  int gx_bit_0 =  gx & 1;
  int gx_bit_1 = (gx & 2) >> 1;
  val = (1-gx_bit_1)*(1-gx_bit_0)*op1(val) + (1-gx_bit_1)*(gx_bit_0)*op2(val) + (gx_bit_1)*(1-gx_bit_0)*op3(val) + (gx_bit_1)*(gx_bit_0)*op4(val);

  g_data[gx] = val;
}

const int N = 32;
int main(){

  float *data;
  cudaMallocManaged(&data, N*sizeof(float));
  for (int i = 0; i < N; i++) data[i] = 1.0f;
  k<<<1,N>>>(data);
  cudaDeviceSynchronize();
  for (int i = 0; i < N; i++) std::cout << data[i] << std::endl;
}
$ nvcc -o t1914 t1914.cu
$ compute-sanitizer ./t1914
========= COMPUTE-SANITIZER
2
3
4
5
2
3
4
5
2
3
4
5
2
3
4
5
2
3
4
5
2
3
4
5
2
3
4
5
2
3
4
5
========= ERROR SUMMARY: 0 errors
$

Answer 2

通过改变每个线程的工作来解决

现有数据布局的最佳解决方案是让每个线程计算 4 个连续值。可以正常工作的线程越少越好。不能正常工作的线程越多。

float* g_data;
int gx = threadIdx.x + blockDim.x * blockIdx.x;
g_data[4 * gx] = op1(g_data[4 * gx]);
g_data[4 * gx + 1] = op2(g_data[4 * gx + 1]);
g_data[4 * gx + 2] = op3(g_data[4 * gx + 2]);
g_data[4 * gx + 3] = op4(g_data[4 * gx + 3]);

如果 g_data 的大小不是 4 的倍数，则在索引操作周围加上 if。如果它始终是 4 的倍数并正确对齐，则将 4 个值加载并存储为 float4 以获得更好的性能。

通过重新排序工作的解决方案

正如我所有关于 float4 的讨论所暗示的那样，您的输入数据似乎是某种形式的 2D 结构，其中每四个元素中就有一个具有相似的功能。也许它是一个结构数组或一个向量数组——换句话说，一个矩阵。

为了解释我的意思，我认为它是一个 Nx4 矩阵。如果将其转置为 4xN 矩阵并对其应用内核，则大多数问题都会消失。因为必须执行相同操作的条目在内存中彼此相邻放置，这使得编写高效内核更容易。像这样的：

float* g_data;
int rows_in_g;
int gx = threadIdx.x + blockDim.x * blockIdx.x;
int gy = threadIdx.y;
float& own_g = g_data[gx + rows_in_g * gy];
switch(gy) {
case 0: own_g = op1(own_g); break;
case 1: own_g = op2(own_g); break;
case 2: own_g = op3(own_g); break;
case 3: own_g = op4(own_g); break;
default: break;
}

以 2D 内核开始，块大小为 x=32，y=4，网格大小为 x=N/32，y=1。

现在你的内核仍然是发散的，但是一个 warp 中的所有线程将执行相同的 case 并访问内存中的连续浮点数。这是你能做到的最好的。当然这一切都取决于您是否可以更改数据布局。