目的
HaskellでGPUを扱う場合にはaccelerateなどがありますが、既存のGPUを扱えるC/C++のコードを使ってGPUとやり取りする場合に、直接CUDAのコードを扱いたい場合があります。
現状では、CUDAのファイルを使うワークフローが煩雑です。
cabalファイルはCUDAをサポートしてないので、CUDAのコードをあらかじめMakefileなどを使って事前にビルドし、cabalファイルにそのライブラリのファイルを追加するという流れになるかと思います。つまりcabalだけでビルドできません。
こちらの例ではMakefileでnvccを使ってCUDAのファイルをビルドしてghcでそのファイルにリンクしていると思います。
今回HaskellのコードにCUDAのコードを埋め込み、cabalファイルに特別な変更を入れないでCUDAのコードを扱えるようにし、cabalだけでビルドできるようにします。
CUDAを埋め込む方法
今回行ったCUDAを埋め込む方法を紹介します。
すでにHaskellのコード中にC言語のコードを埋め込むinline-cパッケージがあります。
こちらは、template-haskellを使って次のように動作します。cabalの変更は不要です。
- Haskellのコードをパースして、C言語のコードを抜き出す。
- HaskellのコードにFFIの
foreign import ccallを使いCの関数を呼び出すコードを追加 - addForeignSourceを使って、GHCにC言語のコードを生成してリンクさせる。
addForeignSourceはC言語の他にC++、Objective C/C++、アセンブラ、オブジェクトファイルを追加することをサポートしています。
今回はオブジェクトファイルを追加する機能を使ってCUDAを呼び出すようにしました。
次のようにしてCUDAを呼び出します。実際はinline-cのコードを流用しています。
- Haskellのコードをパースして、CUDAのコードを抜き出す。
- CUDAのコードをnvccでコンパイルしてオブジェクトファイルにする。
- HaskellのコードにFFIの
foreign import ccallを使いC++の関数を呼び出すコードを追加（CUDAとC++の関数は同じ命名規約） - addForeignSourceを使って、GHCにCUDAのオブジェクトファイルをリンクさせる。
ソースコードはこちらです。
記述例
CUDAをHaskellに埋め込む例は次のようなものです。完全な例はこちらです。
-- C言語でなくてCUDAを使うように設定
import qualified Language.C.Inline.Cuda as C
C.context $ C.cudaCtx
-- CUDAのカーネルを書く。
[C.emitBlock|
__global__ void
vectorAdd(const float *A, const float *B, float *C, int numElements)
{
...
}
|]
-- CUDAのカーネルを呼び出す。
vectorAdd :: Ptr C.CFloat -> Ptr C.CFloat -> Ptr C.CFloat -> CInt -> IO ()
vectorAdd d_A d_B d_C cNumElements =
[C.block| void {
const int threadsPerBlock = 256;
const int blocksPerGrid =($(int cNumElements) + threadsPerBlock - 1) / threadsPerBlock;
vectorAdd<<<blocksPerGrid, threadsPerBlock>>>($(float* d_A), $(float* d_B), $(float* d_C), $(int cNumElements));
} |]
まとめと課題
HaskellとCUDAの連携を既存のワークフローを壊さずにcabalだけで扱えるライブラリinline-c-cudaを作成しました。（nvccといったCUDAのコンパイラは別途インストールが必要です。）
今後の課題は、 CUDAのカーネル関数のFusionだと思っています。最近のGPUはメモリのバンド幅がネックになってきているらしく、複数のカーネル関数の一つの関数にしてメモリアクセスを減らし、パフォーマンスを上げる最適化がpytorchやCuPyでは行われています。CUDAを直接使わずにpythonでカーネル関数を使って行われています。
HaskellのStream Fusionで何とかなるといいのですが。