Cuda C を用いた並列リダクションの実装と高速化

最近, cuda C を使う機会が多くて, リファレンスを探すのに苦労したので, そのメモです。
この記事は,

という流れです。

atomicAdd 関数の自作

複数のスレッドである変数を読み取り, 四則演算などをしてその値を書き込むときに使用します。
https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
倍精度浮動小数点数同士の和も, Compute Capability 6.0 以上ではatomicAdd が使えるらしい。そうでない場合は以下のように組む。

__device__ double atomicAdd_double(double* address, double val)
{
    unsigned long long int* address_as_ull =(unsigned long long int*)address;
    unsigned long long int old = *address_as_ull, assumed;

    do {
        assumed = old;
        old = atomicCAS(address_as_ull, assumed,
                        __double_as_longlong(val + __longlong_as_double(assumed)));
    } while (assumed != old);

    return __longlong_as_double(old);
}

並列リダクションの実装

https://devblogs.nvidia.com/faster-parallel-reductions-kepler/
をみた感じだと, CUB Library を使用するか warp+atomicAdd を使用すると良さそう。
version が古いと, double 型に対する shuffle_down 操作も書く必要があるので, まずはそのバージョン。

__device__ double __shfl_down_double(double var, unsigned int srcLane, int width = 32) {
	int2 a = *reinterpret_cast<int2*>(&var);
	a.x = __shfl_down_sync(0xffffffff, a.x, srcLane, width); // or a.x = __shfl_down(a.x, srcLane, width);
	a.y = __shfl_down_sync(0xffffffff, a.y, srcLane, width); // or a.y = __shfl_down(a.y, srcLane, width); 
	return *reinterpret_cast<double*>(&a);
}

__device__ double warpReduceSum_double(double val) {
	for (int offset = warpSize / 2; offset > 0; offset /= 2)
		val += __shfl_down_double(val, offset);
	return val;
}

__global__ void deviceReduceWarpAtomicKernel(double *in, double* out, int N) {
	double sum = double(0.0);
	for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; i += blockDim.x * gridDim.x) {
		sum += in[i];
	}
	sum = warpReduceSum_double(sum);
	if ((threadIdx.x & (warpSize - 1)) == 0)
		atomicAdd_double(out, sum);
}

deviceReduceWarpAtomicKernel のfor 文では, grid stride loop が用いられています。
CUDA Pro Tip: Write Flexible Kernels with Grid-Stride Loops | NVIDIA Developer Blog

Programming Guide :: CUDA Toolkit Documentation
を見てみると, __shfl_down_sync は int, unsigned int, long, unsigned long, long long, unsigned long long, float, doubleに対応しているらしい。また
Using CUDA Warp-Level Primitives | NVIDIA Developer Blog
warp に関する情報が記載されています。
CUDA9.0 以降だと,

__device__ double warpReduceSum_double(double val) {
	for (int offset = (warpSize >> 1); offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffff, val, offset);
	return val;
}

__global__ void deviceReduceWarpAtomicKernel(double *in, double* out, int N) {
	double sum = double(0.0);
	for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; i += blockDim.x * gridDim.x) {
		sum += in[i];
	}
	sum = warpReduceSum_double(sum);
	if ((threadIdx.x & (warpSize - 1)) == 0)
		atomicAdd_double(out, sum);
}

で, OK。

ベクトル化による高速化(あまり効果はなかった)

さらに, 先ほどのコードをベクトル化することにより高速化してみる。
CUDA Pro Tip: Increase Performance with Vectorized Memory Access
が参考になります。

__global__ void deviceReduceWarpAtomicVectorized2Kernel(double *in, double* out, unsigned int N) {
    double sum = double(0.0);
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int half_N = N >> 1;
    double2 tmp;
    for (unsigned int i = idx; i < half_N; i += blockDim.x * gridDim.x) {
        tmp = reinterpret_cast<double2*>(in)[i];
        sum += tmp.x + tmp.y;
    }

    sum = warpReduceSum_double(sum);
    if ((threadIdx.x & (warpSize - 1)) == 0) atomicAdd_double(out, sum);
}

以下では, N = 2**26 として, 実行します。
nvprof を用いて大雑把に実行時間の比較:
deviceReduceWarpAtomicKernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024>>> 798.97
<<< 131072, 512>>> 746.29
<<< 262144, 256>>> 573.90
<<< 524288, 128>>> 577.57
<<< 131072, 256>>> 286.78
<<< 65536, 256>>> 143.56
<<< 32768, 256>>> 72.505
<<< 16384, 256>>> 37.565
<<< 8192, 256>>> 20.420
<<< 4096, 256>>> 12.038
<<< 2048, 256>>> 8.5510
<<< 1024, 256>>> 8.0172
<<< 512, 256>>> 7.6053
<<< 256, 256>>> 7.5142
<<< 128, 256>>> 7.4906

deviceReduceWarpAtomicVectorized2Kernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024 >>> 408.98
<<< 131072, 512 >>> 385.48
<<< 262144, 256 >>> 298.33
<<< 524288, 128 >>> 300.19
<<< 131072, 256 >>> 282.86
<<< 65536, 256 >>> 141.58
<<< 32768, 256 >>> 71.043
<<< 2048, 256 >>> 8.2606
<<< 1024, 256 >>> 7.5353
<<< 512, 256 >>> 7.4422
<<< 256, 256 >>> 7.4634
<<< 128, 256 >>> 7.4942

同様にして,

__global__ void deviceReduceWarpAtomicVectorized4Kernel(double *in, double* out, unsigned int N) {
    double sum = double(0.0);
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int loop_N = N >> 2;
    double4 tmp;
    for (unsigned int i = idx; i < loop_N; i += blockDim.x * gridDim.x) {
        tmp = reinterpret_cast<double4*>(in)[i];
        sum += tmp.x + tmp.y + tmp.z + tmp.w;
    }

    sum = warpReduceSum_double(sum);
    if ((threadIdx.x & (warpSize - 1)) == 0) atomicAdd_double(out, sum);
}

deviceReduceWarpAtomicVectorized4Kernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024>>> 213.49
<<< 131072, 512>>> 180.72
<<< 262144, 256>>> 166.55
<<< 524288, 128>>> 167.64
<<< 65536, 256>>> 142.91
<<< 2048, 256>>> 8.2487
<<< 1024, 256>>> 7.7647
<<< 512, 256>>> 7.9956

ベクトル化はやりすぎても良くないみたいで, 効果はいまいちみたいです… ベクトル化を使う場合にdouble2 までというのは事実っぽいですね。それ以上に, ブロック当たりの thread 数やグリッド当たりのブロック数のチューニングをする方が, 性能がかなり変わりますね!

複素数の場合

同様のことが複素数に対してもできます。

__global__ void deviceReduceWarpAtomicKernel(cuDoubleComplex *in, cuDoubleComplex* out, unsigned int N) {
	cuDoubleComplex sum = make_cuDoubleComplex(0.0, 0.0);
	for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; i += blockDim.x * gridDim.x) {
		sum = cuCadd(sum, in[i]);
	}
	sum.x = warpReduceSum_double(sum.x);
	sum.y = warpReduceSum_double(sum.y);
	if ((threadIdx.x & (warpSize - 1)) == 0){
		atomicAdd_double(&out[0].x, sum.x);
		atomicAdd_double(&out[0].y, sum.y);
    }
}

__global__ void deviceReduceWarpAtomicVectorized2Kernel(cuDoubleComplex *in, cuDoubleComplex* out, unsigned int N) {
	cuDoubleComplex sum = make_cuDoubleComplex(0.0, 0.0);
    double4 tmp;
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int half_N = N >> 1;
	for (unsigned int i = idx; i < half_N; i += blockDim.x * gridDim.x) {
        tmp = reinterpret_cast<double4*>(in)[i];
		sum.x += tmp.x + tmp.z;
        sum.y += tmp.y + tmp.w;
	}
	sum.x = warpReduceSum_double(sum.x);
	sum.y = warpReduceSum_double(sum.y);
	if ((threadIdx.x & (warpSize - 1)) == 0){
		atomicAdd_double(&out[0].x, sum.x);
		atomicAdd_double(&out[0].y, sum.y);
    }
}

複素数の場合の実行時間の比較:
deviceReduceWarpAtomicKernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024>>> 831.83
<<< 131072, 512>>> 796.84
<<< 262144, 256>>> 574.31
<<< 524288, 128>>> 577.55
<<< 256, 256>>> 15.069
<<< 128, 256>>> 14.850

deviceReduceWarpAtomicVectorized2Kernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024 >>> 438.12
<<< 131072, 512 >>> 412.84
<<< 262144, 256 >>> 319.97
<<< 524288, 128 >>> 321.76
<<< 131072, 256 >>> 290.40
<<< 4096, 256 >>> 17.290
<<< 2048, 256 >>> 15.544
<<< 1024, 256 >>> 15.709
<<< 256, 256 >>> 15.889
<<< 128, 256 >>> 16.083

ブロックあたりのスレッド数は 256 で良さそうなので, 以下ではスレッド数を256 とします。

高速化の結果

さらに, 同様の方針で自作をして実行時間を計測したものをまとめると以下の表のようになります。

関数名 実数の場合の実行時間 (ms) 複素数の場合の実行時間 (ms)
deviceReduceWarpAtomicKernel (事実上 grid stride なし) 798.97 831.83
deviceReduceWarpAtomicKernel 7.4906 14.850
deviceReduceWarpAtomicVectorized2Kernel 7.4422 15.544
deviceReduceWarpAtomicVectorized4Kernel 7.7647 24.040

並列リダクションを行う場合, atomic 操作がボトルネックになりやすいので, それを減らすようにすると速くなるというぐらいです。