Grid-stride loop - 東方算程譚

NVIDIAの開発系BlogになるほどなーなTips: Grid-stride loop を見つけたので軽く解説。 (元ネタはコチラ)

毎度おなじみ SAXPY : ベクトルの積和演算 Y[i] = a * X[i] + Y[i] (i = 0,1 ... n-1) これをCPUでフツーに行うとき、iについてくるくる回すloopで実装しますわな。

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

かたやこいつをGPU(CUDA)でやるときの常套手段は "iによるn回のloop" を "n個のスレッド"　に置き換えることでコードからloopを取っ払います。

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

void saxpy(int n, float a, const float* x, float* y) {
  kernel_saxpy<<<(n+255)/256, 256>>>(n, a, x, y);
}

Grid-stride loopは両者の中間みたいなやりかたで、for-loopを温存します。

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

このとき iの初期値はスレッドの通し番号、loopごとの増分はスレッドの総数すなわち gridDim.x * blockDim.x です。するってーとスレッド総数が Tのとき、

 スレッド#0 : i = 0, T  , 2T ...  
 スレッド#1 : i = 1, T+1, 2T+1 ...  
 スレッド#2 : i = 2, T+2, 2T+2 ...  
 (以下同文)

に対して y[i] を計算してくれる、と。

こうしておくとなにがオイシイかというと、kernel-call時のスレッド数を自由に指定できます。grid * blockがnより小さくてもちゃんと長さnのベクトルについて計算してくれる。

なのでたとえばGPUが持ってるプロセッサ数に応じてgrid/block数をチューニングできます。

void saxpy(int n, float a, const float* x, float* y) {
  int numSMs;
  cudaDeviceGetAttribute(&numSMs, cudaDevAttrMultiProcessorCount, 0);
  // SMあたり32*256=8192個のスレッドを起こして計算する
  kernel_saxpy<<<numSMs*32, 256>>>(n, a, x, y);
}

極端な話 kernel_saxpy<<<1, 1>>>(n, a, x, y); でも問題ありません。このときスレッド一個で処理しますから多数のスレッドが絡む問題を起こさないので本来得られるべき結果が出てくるわけで、デバッグ時の検証に重宝しますね。加えて kernel内で'prinf'したとき出力がお行儀よく並びますからデバッグが楽ですよね。

おためしに書いてたコードがこちらになります:

#include <cuda_runtime.h>
#include <device_launch_parameters.h>

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

#include <iostream>
#include <random>
#include <algorithm>
#include <ctime>

int main() {
  using namespace std;

  const int N = 100000;
  
  float* x;
  float* y0;
  float* y;
  float  a;

  cudaMallocManaged(&x, N*sizeof(float));
  cudaMallocManaged(&y, N*sizeof(float));
  cudaMallocManaged(&y0, N*sizeof(float));

  mt19937 gen(time(nullptr));
  uniform_real_distribution<float> dist;
  auto rand = [&]() { return dist(gen); };
  generate_n(x, N, rand);
  generate_n(y, N, rand);
  copy(y, y+N, y0);
  a = rand();

  // 単一スレッドで計算する
  kernel_saxpy<<<1,1>>>(N, a, x, y0);

  int numSMs;
  cudaDeviceGetAttribute(&numSMs, cudaDevAttrMultiProcessorCount, 0);
  cout << numSMs << " SMs" << endl;
  // SMあたり32*256=8192個のスレッドを起こして計算する
  kernel_saxpy<<<numSMs*32, 256>>>(N, a, x, y);
  cudaDeviceSynchronize();

  if ( equal(y0, y0+N, y) ) {
    cout << "ok." << endl;
  }

  cudaFree(x);
  cudaFree(y);
  cudaFree(y0);
  cudaDeviceReset();

}