CUDAでスレッドやブロックの統一されたインデックスや全個数などを得る関数群

最近ブログをサボりがちなので、更新のためにCUDAで自分が普段使ってるソースを晒す。

#ifndef __KERNEL_UTILITY_H__
#define __KERNEL_UTILITY_H__

#include <cutil.h>

// 一次元配列を二次元配列として扱うとき、二次元配列のarray[x][y]が一次元配列のどこにあたるかを計算する関数
// ただし、二次元配列を一次元配列にマッピングするとき、どのような配置にするかでコアレッシングが発生する/発生しない
// などの速度の違いが発生するため、この実装でいいのかはわからない。
// もしかすると場合によっては、 X_LENGTH * y + x の方が高速な可能性がある。
__device__ __host__ int getSerialIndex(const int Y_LENGTH, const int x, const int y) {
	return Y_LENGTH * x + y;
}

// この関数を実行したブロックに含まれるスレッドの総数
__device__ int getThreadCountPerBlock() {
	return blockDim.x * blockDim.y * blockDim.z;
}

// この関数を実行したグリッドに含まれるブロックの総数
__device__ int getBlockCountPerGrid() {
	return gridDim.x * gridDim.y * gridDim.z;
}

// スレッドの固有なインデックス(ブロック内の)を返す
__device__ int getThreadIndexPerBlock() {
	return ((threadIdx.z * blockDim.y) + threadIdx.y) * blockDim.x + threadIdx.x;
}

// ブロックの固有なインデックス(グリッド内の)を返す
__device__ int getBlockIndexPerGrid() {
	return ((blockIdx.z * gridDim.y) + blockIdx.y) * gridDim.x + blockIdx.x;
}

// グリッド内に含まれるスレッドすべてに、それぞれ固有のインデックスを割り当てる関数
__device__ int getGlobalUniqueIndexPerGrid() {
	return getBlockIndexPerGrid() * getThreadCountPerBlock() + getThreadIndexPerBlock();
}

// グリッド内に含まれるスレッドの総数。つまり、ブロックの総数*ブロック内に含まれるスレッドの総数
__device__ int getOverallThreadCountPerGrid() {
	return getBlockCountPerGrid() * getThreadCountPerBlock();
}

#endif /* __KERNEL_UTILITY_H__ */

こんなものはぶっちゃけだれでも書けるし、事実これと同等のコードを毎回書く人も多いだろう。
ただ、ある程度使用されて信頼性の高まった関数を使用するメリットは非常に高い。つまらないtypoや変数の指定ミス*1がなくなるだけで余計なリスクは大幅に減る。CやC++系統にどっぷり浸かった人は1から10まで自分で書かないと納得がいかない人も多いが*2、ミドルコードでも使えるものはどんどん使えばいいと思う。

*1:私はgridDimとblockDimを非常に混同しやすい

*2:実際そういった態度が必要な分野もある