CUDA Warpシャッフル命令のエミュレーション - syghの新フラグメント置き場

今更ですがせっかくCompute Capability 3.0対応のKepler世代グラフィックスカードを手に入れたので、CUDAのWarpシャッフル命令の動作テストを兼ねて、代替機能をエミュレートする関数を書いてみました。
Visual Studio 2012、CUDA 6.5、GeForce GTX 770で動作確認済み。

#define WARP_SIZE (32)

typedef unsigned int uint;

// HLSL の SV_GroupIndex 相当。
static __device__ __forceinline__ int GetLocalThreadGroupFlatIndex()
{
  return threadIdx.x + blockDim.x * threadIdx.y + blockDim.x * blockDim.y * threadIdx.z;
}

static __device__ __forceinline__ int GetGlobalThreadFlatIndex()
{
  const int blockId = blockIdx.x + gridDim.x * blockIdx.y + gridDim.x * gridDim.y * blockIdx.z;
  return blockId * blockDim.x * blockDim.y * blockDim.z + GetLocalThreadGroupFlatIndex();
}

#pragma region // Emulation Functions //

// 書き込み後の同期だけでなく、読み込み後の同期も必要。でないと安全に共有メモリ上の古いデータを上書き／破棄できない。

template<typename T> static __device__ T EmulateShuffle(T x, int srcLane)
{
  extern __shared__ T sharedMem[];
  const int localIndex = GetLocalThreadGroupFlatIndex();
  const int laneBlockId = localIndex / WARP_SIZE;
  const int laneId = srcLane & (WARP_SIZE - 1);
  sharedMem[localIndex] = x;
  __syncthreads();
  const T y = sharedMem[laneBlockId * WARP_SIZE + laneId];
  __syncthreads();
  return y;
}

template<typename T> static __device__ T EmulateShuffleUp(T x, uint delta)
{
  extern __shared__ T sharedMem[];
  const int localIndex = GetLocalThreadGroupFlatIndex();
  const int laneBlockId = localIndex / WARP_SIZE;
  const int laneId = localIndex & (WARP_SIZE - 1);
  const int shiftAmt = laneId - static_cast<int>(delta);
  sharedMem[localIndex] = x;
  __syncthreads();
  const T y = (0 <= shiftAmt) ? sharedMem[laneBlockId * WARP_SIZE + shiftAmt] : sharedMem[localIndex];
  __syncthreads();
  return y;
}

template<typename T> static __device__ T EmulateShuffleDown(T x, uint delta)
{
  extern __shared__ T sharedMem[];
  const int localIndex = GetLocalThreadGroupFlatIndex();
  const int laneBlockId = localIndex / WARP_SIZE;
  const int laneId = localIndex & (WARP_SIZE - 1);
  const int shiftAmt = laneId + static_cast<int>(delta);
  sharedMem[localIndex] = x;
  __syncthreads();
  const T y = (shiftAmt <= (WARP_SIZE - 1)) ? sharedMem[laneBlockId * WARP_SIZE + shiftAmt] : sharedMem[localIndex];
  __syncthreads();
  return y;
}

template<typename T> static __device__ T EmulateShuffleXor(T x, int laneMask)
{
  extern __shared__ T sharedMem[];
  const int localIndex = GetLocalThreadGroupFlatIndex();
  const int laneBlockId = localIndex / WARP_SIZE;
  const int laneId = localIndex & (WARP_SIZE - 1);
  sharedMem[localIndex] = x;
  __syncthreads();
  const T y = sharedMem[laneBlockId * WARP_SIZE + ((laneId ^ laneMask) & (WARP_SIZE - 1))];
  __syncthreads();
  return y;
}

#pragma endregion

エミュレート関数はそれぞれ、下記のCUDA組み込み関数に相当します。当然組み込み関数を使ったほうが手軽だし高速ですが、CC 3.0以降のデバイスでないと使えません。

// 指定したレーン (Warp 内のスレッド番号) の var の値を受け取ります。
T __shfl(T var, int srcLane, int width = warpSize);

// 指定した数 delta だけ離れたスレッドの var の値を受け取ります。
// up は指定した数だけ上の番号のスレッドから、down は下の番号のスレッドから受け取ります。
// Warp の外に出てしまった場合は自身のスレッドでの値が返ってきます。
T __shfl_up(T var, unsigned int delta, int width = warpSize);
T __shfl_down(T var, unsigned int delta, int width = warpSize);

// 自分のレーン番号と指定した laneMask を XOR した結果の番号のスレッドから受け取ります。
// laneMask ごとにレジスタの中身を交換することができます。
T __shfl_xor(T var, int laneMask, int width = warpSize);

CUDAの場合、共有メモリのサイズ指定はコンパイル時だけでなく、下記のように（ハードウェア仕様上限が許す限り）実行時に指定することもできます。今回のエミュレーション関数は共有メモリにexternを指定することで、実行時にサイズ指定する方法を選びました。

const dim3 gridSize(16, 4, 4);
const dim3 blockSize(64, 8, 2);
const size_t elemCount = gridSize.x * gridSize.y * gridSize.z * blockSize.x * blockSize.y * blockSize.z;
const size_t sharedMemSize = blockSize.x * blockSize.y * blockSize.z * sizeof(TargetType);
…
kernelFunc<<<gridSize, blockSize, sharedMemSize>>>(devOutArray, devInArray);

一応動作確認テストはしましたが、完全に同一仕様でエミュレートできているかどうかは未保証です。ご利用の際は自己責任でどうぞ。

CUDAはNVIDIA専用という最大のネックがあるので、個人的にDirectComputeに比べてあまり好きではないのですが、CUDA自体は非常に効率よく最高水準のGPGPUプログラムを記述できる、地球上でもっとも優れた超並列プログラミング環境であることは確かです。C++erとして特に気に入っているのは、Visual C++などのホストコンパイラーにコンパイルさせるホスト関数だけでなく、NVCCコンパイラーにコンパイルさせるデバイス関数でもクラス、テンプレートやC++11 Featureがほとんどそのまま使える点です（CUDA 7ではC++11対応が強化されるそうです）。HLSLでもクラスは使えますが、CUDAのプログラマビリティはHLSLを遥かに凌駕します。GLSLコンピュートシェーダーやOpenCL-Cごときではまったく勝負にならないレベルです。
なおネイティブC++から使えるGPGPU言語拡張としては、Visual C++ 2012以降のC++ AMPも強力です。WindowsやXbox Oneのほか、すでにLinux向けにもオープンソース実装されているので、OpenCLの代わりにC++ AMPを使ってみてもよいでしょう。

ちなみに自分はPC用グラフィックスカードに関しては、G7xからG8x、Fermi、Keplerと乗り換えてきましたが、ずっとNVIDIAオーナーです。AMDはハードウェアに関しては良いモノを作っていることは確かなんですが（ゲーム向けも業務向けもHPC向けも性能に関してはAMDのほうが上）、ドライバーの完成度がお話にならないレベルであったり、ソフトウェアを軽視しているメーカーだと感じます。また、APUを除き、現行のグラフィックスカード単体での省電力性能は、ハイエンドでもNVIDIAに分があります。Keplerを使い倒したら、いずれMaxwellへ、さらにその後はPascalへ移行する予定ですが、AMDが今の劣悪な状況を改善しようとしないかぎり、今後も個人でAMDグラフィックスを購入することはないでしょう（逆に言うとドライバー／ソフトウェア面さえ改善されれば、オープン標準規格のOpenCLやC++ AMPを推進するAMDはNVIDIAよりも良い選択肢になり得ます）。