cuda(nvcc)でpythonのモジュールを作ってみた

2015-11-18 00:06

ここのところcudaが楽しくてしょうがないのですが、でもC言語で高いレイヤーの処理はやりたくないですよね。画像をCで触るとか、考えたくもない。そんなわけでふつう使うのはpycudaあたりだと思うのですが、あえてPython/C APIでモジュールを作ってみることにしました。

とりあえずソースを書きます。

#include <python3.4/Python.h>  // 環境によってはバージョンが違ったり、<python.h>だけで良かったりするかも。


static PyObject* tea_device_count(PyObject *self){
    int n;
    cudaGetDeviceCount(&n);
    return PyLong_FromLong(n);
}


static PyMethodDef TeaMethods[] = {
    {"device_count", (PyCFunction)tea_device_count, METH_NOARGS, "get device count."},
    {NULL, NULL, 0, NULL}
};


static struct PyModuleDef teamodule = {
    PyModuleDef_HEAD_INIT,
    "tea",
    "this is tea\n",
    -1,
    TeaMethods
};


PyMODINIT_FUNC PyInit_tea(){
    return PyModule_Create(&teamodule);
}

こんな感じ。teaってモジュールに、device_countってGPUの数を数える関数が一つあるだけです。

コンパイルするときはこんな。

$ nvcc -shared -Xcompiler -fPIC tea.cu -o tea.so

-fPICってやつが-Xcompilerの前に来るとエラーになるので要注意です。あとは普通に？

使い方は完全に普通のpythonモジュールと同じで、以下のような感じになります。

>>> import tea
>>> tea.device_count()
1

お使いのGPUの数が表示されると思います。2とか3とか出たらリッチな感じですね。

計算などはまだ試していませんが、まあここまで出来て出来ないことはないかと。素直にpycudaを使えって感じもしますが、これはこれで悪くない気がします。

参考：