NVRTCで見るCUDAのカーネル関数

NVRTCが何か

NVRTC(NVIDIA RunTime Compilation)はchar配列に収まっているカーネル関数のソースコードをコンパイルして呼び出すためのAPIです．
一度使ってみると，カーネル関数のコンパイルコードを書くこととなるため，CUDAのカーネル関数がどう実行されるか理解するのに役立ちます．
NVRTCはcupyのelementwiseみたいなことに使えます．

NVRTCの処理の流れ

1. CUDA/Cのカーネル関数コードからPTXコードを生成
2. PTXコードからCUfunctionを生成
3. 引数を設定しCUfunctionを立ち上げ

みたいな感じです．

コード例

(結構昔に書いたコードで汚くてごめんなさいm(_ _)m)

1. PTXコードを生成

1. nvrtcCreateProgram関数でカーネル関数の文字列char配列やincludeさせて展開させたいファイルなどを登録します．
2. nvrtcCompileProgram関数でコンパイルします． この際コンパイルオプションを渡します．
3. nvrtcGetPTXSize関数でPTXコードを文字列として取得します．

2で指定するコンパイルオプションですが，例えば--arch=sm_70のようなものです． これらのコンパイルオプションはPTXを生成するために使われ，バイナリ生成に必要なわけではないようです． (仮想アーキテクチャ指定なので__CUDA_ARCH__の値を設定したりする感じですよね．)

extern "C"

がないとマングルされて関数名が解決できなくなるので注意が必要です．　

2. PTXコードからCUfunctionを生成

CUfunctionを作成するために必要なものはカーネルのPTXコードのみで，あとは動かすデバイスに合わせてCUfunctionを作成するようです． そもそもnvccでコンパイルしたCUDAプログラムは内部に"PTX"を含んでおり，実行時に実行GPUに適したカーネル関数イメージがない場合は実行時にコンパイルしてリンクさせます．

PTXから自動で適した形式の関数実体にコンパイルしてくれるため，こちらがコンパイルオプションの指定する必要もないようです．

3. 引数を設定しCUfunctionを立ち上げ

引数の"ポインタ"を配列にして渡します．
ポインタの配列を用いる利点はいろいろありそうですが，void*の配列として引数に取れるのでcuLaunchKernel関数の設計が楽になりますよね．
その他にもdim3類やShared Memoryのサイズ制限やStreamなどを渡します．

終わりに

コンパイルオプションが何を生成する際に使われるのかなど，NVRTCのコードを書くことで得られる知見も多々あります．(nvccのDocumentをちゃんと読めば書いてあるのかもしれないですけど．)
ぜひぜひいろいろ書いて遊んでみてください．

参考

• CUDA Toolkit Documentation - NVRTC