Write_Flexible_Kernels_with_Grid-Stride_Loops

Write Flexible Kernels with Grid-Stride Loops

2020 Jan 28th CookieLau

Source: https://devblogs.nvidia.com/cuda-pro-tip-write-flexible-kernels-grid-stride-loops/

举一个向量运算的简单例子 SAXPY：

void saxpy(int n, float a, float *x, float *y)
{
    for (int i = 0; i < n; ++i)
        y[i] = a * x[i] + y[i];
}

通常的 cuda 加速就是分配足够多的线程，然后合法的线程去执行其核心代码：

__global__
void saxpy(int n, float a, float *x, float *y)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) 
        y[i] = a * x[i] + y[i];
}

我们通常称这种 cuda 加速的方法叫做 monolithic kernel，因为其需要创建的线程会随着需要运算的数据量的增大而增大，比如下面如果我的向量长度是 1M，则需要创建 1M 个线程：

// Perform SAXPY on 1M elements
saxpy<<<4096,256>>>(1<<20, 2.0, x, y);

显然很不美观～

这里我们推荐另一种方法： grid-stride loop 即以 grid 的大小为布长的小循环运算：

__global__
void saxpy(int n, float a, float *x, float *y)
{
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; // 也是从自身thread_id开始
         i < n; 
         i += blockDim.x * gridDim.x) 
         // 不过每一个线程不止负责自己，还负责每个 grid 的自己位置上的线程
      {
          y[i] = a * x[i] + y[i];
      }
}

举个例子，如果一个 grid 中有 100 个线程，则 thread_id=0 的线程会负责 0，100，200，etc

我们看到 monolithic 中的例子不过是将 grid-Stride 中的循环展开而已，所以时间复杂度两者来说应该是相近的，测试结果也确实相近：

向量相加，大小 1<<25

Device ID: 0    Number of SMs: 80
numberOfBlocks: 32*numberOfSMs
threadsPerBlock: 256
grid-Stride time cost:525.279999 <<< numberOfBlocks, threadsPerBlocks >>>
monolithic time cost:502.784014 <<< 1<<17, 256 >>>

虽然 monolithic 快一点点，但是 grid-Stride 可以带来其他的好处如下： 1. Scalability and thread reuse: 在上面的例子也可以看出来，monolithic 在调用核函数的时候启动的 gridSize 和 blockSize 的大小是硬编码的，如果后续数据量发生了变化则需要修改硬编码，而 grid-Stride 则是只由 numberOfSMs 决定，不会由数据量的大小发生改变，更加的 Scalable。此外，monolithic 的方式要一次性分配所有的线程，那么如果极端情况下，gridSize*blockSize 无法一次性分配那么多，超过了 GPU 的能力，则无法执行核函数。 2. 对 grid-Stride 的调参也非常简便，只需要调节 numberOfBlock 一个参数即可。 3. grid-Stride 实现了对 thread 的重用(reuse)，节省了单个线程的创建和销毁的时间开销还有其他的自定义的在每个线程的前后的动作(可有可无，如 thread-private and shared data initialization)。 4. Debugging 友好: 在 grid-Stride 方式下，只需要将核函数调用参数改成 <<>> 就可以实现 serial processing，通过串行来调试bug。 5. Portability and Readability: 大概就是更符合原来的 sequential 的感觉，给其他用户更方便阅读。