PythonでC拡張を書く - pyhaya’s diary

Pythonは速度で見ると早いとは言えない言語です。しかし、C言語による拡張を書くことができて、それにより速度を大幅に上昇させることができます。よく言語の速さを比較するのに使われるフィボナッチ数列を使ってピュアPythonとC拡張の速度の比較を行います。

PythonのC拡張ではよくCythonが話題に上ります。これはPythonのような文法で簡単にC拡張を書くことができるため人気があります。しかし、この記事では純粋にC言語から出発してPythonへコードを移植します。

この記事は「エキスパートPythonプログラミング」を参考にしています。

エキスパートPythonプログラミング改訂2版

作者: Michal Jaworski,Tarek Ziade,稲田直哉,芝田将,渋川よしき,清水川貴之,森本哲也
出版社/メーカー: KADOKAWA
発売日: 2018/02/26
メディア: 単行本
この商品を含むブログを見る

開発環境
必要となるもの
Pythonでの実装
Cでの実装
Pythonで使える形にする
拡張モジュールをコンパイルする
使ってみる
さらに最適化してみる

開発環境

Windows10
Python 3.6.5
Anaconda

必要となるもの

C言語のコンパイルが必要になるので、gccやVisual Studio等が必要になります。私は、コンパイルはWSLでやっておりますのでgccを使っています。また、Pythonで使える形に書くためにPython.hというヘッダファイルが必要になります。私のようなAnacondaでPythonをインストールしている場合にはC:/Users/(ユーザー名)/Anaconda3/include中にあります。

Pythonでの実装

高速化処理は行わずに純粋に再帰だけで書いてみます。

python

def fib(n):
    if n < 2:
        return 1
    return fib(n-1) + fib(n-2)

Cでの実装

C言語でも実装は似たような感じになります。Python実装と名前がかぶらないように名前はfibonacciに変えています。

C言語

int fibonacci(int n){
    if (n < 2){
        return 1;
    }else{
        return fibonacci(n-1) + fibonacci(n-2);
    }
}

Pythonで使える形にする

上のC言語実装だけではPythonで使えません。Pythonで使えるように下のようにコードを付け加えます。

fibonacci.c

#include <Python.h>

int fibonacci(int n){
    if (n < 2){
        return 1;
    }else{
        return fibonacci(n-1) + fibonacci(n-2);
    }
}

static PyObject* fibonacci_py(PyObject* self, PyObject* args){
    PyObject *result = NULL;
    long n;
    if (PyArg_ParseTuple(args, "l", &n)){
        if( n < 0 ){
            PyErr_SetString(PyExc_ValueError, "n must not be less than 0");
        }else{
            result = Py_BuildValue("L", fibonacci((unsigned int)n));
        }
    }

    return result;
}
static char fibonacci_docs[] = "fibonacci(n): Return nth Fibonacci sequence number computed recuesive\n"; 
                                                         
static PyMethodDef fibonacci_module_methods[] = {
    {"fibonacci", (PyCFunction)fibonacci_py,
        METH_VARARGS, fibonacci_docs},
    {NULL, NULL, 0, NULL} 
};

static struct PyModuleDef fibonacci_module_definition = {
    PyModuleDef_HEAD_INIT,
    "fibonacci",
    "Extension module that provides fibonacci sequence function",
    -1,
    fibonacci_module_methods
};

PyMODINIT_FUNC PyInit_fibonacci(void){
    Py_Initialize();
    return PyModule_Create(&fibonacci_module_definition);
}

ずいぶん長いコードになりました。このコードのほとんどはボイラープレートコードです。一つずつ見ていきます。

`fibonacci_py`

static PyObject* fibonacci_py(PyObject* self, PyObject* args){
    PyObject *result = NULL;
    long n;
    if (PyArg_ParseTuple(args, "l", &n)){
        if( n < 0 ){
            PyErr_SetString(PyExc_ValueError, "n must not be less than 0");
        }else{
            result = Py_BuildValue("L", fibonacci((unsigned int)n));
        }
    }

    return result;
}

このコードは、C言語の関数を、Pythonで扱えるオブジェクトを返すようにするためのコードです。Python/C APIはPyObjectという型をこのために用意していて、すべての関数はこの型のポインタを返す必要があります。

PyObject* argsが関数の受け取る引数を含むタプルへのポインタになっています。nという変数を用意しておいてPyArg_ParseTupleでnに入れています。"l"というのは引数がlong型であることを期待していることを示しています。

最後に、fibonacci((unsigned int)n)で数列を計算して、それをPythonで使えるオブジェクトに変換します。変換はPy_BuildValueが行います。

`fibonacci_docs[]`

ドキュメントです。

`fibonacci_module_methods[]`

static PyMethodDef fibonacci_module_methods[] = {
    {"fibonacci", (PyCFunction)fibonacci_py,
        METH_VARARGS, fibonacci_docs},
    {NULL, NULL, 0, NULL} 
};

この配列は、モジュールが提供する関数やメソッドを定義します。配列は以下の要素を含みます。

関数名
関数のC実装へのポインタ
呼び出し規約・束縛条件を含むフラグ
docstring文字列へのポインタ

配列の最後に入っているのは番兵です。C実装へのポインタは、fibonacci_pyのPyCFunctionへのキャストです。この関数の呼び出し規約がMETH_VARAGSで決められています。呼び出し規約にはいくつか選択肢があります。

規約	説明
`METH_VARARGS`	パラメータとして引数リストのみを受け取る
`METH_KWARDS`	キーワード引数を利用できる
`METH_NOARGS`	引数無し

`fibonacci_module_definition`

static struct PyModuleDef fibonacci_module_definition = {
    PyModuleDef_HEAD_INIT,
    "fibonacci",
    "Extension module that provides fibonacci sequence function",
    -1,
    fibonacci_module_methods
};

モジュール全体を定義する構造体です。最初の要素は必ずPyModuleDef_HEAD_INITを使います。第二要素はモジュール名です。第三要素はモジュールのdocstringへのポインタ、第三要素はモジュールの状態を保持するために確保されるメモリの大きさを表しています。これはほとんどの場合には-1で大丈夫で、複数のサブインタープリタや複数段階での初期化が必要な時に使います。第四要素は関数をPyModuleDefで定義した配列へのポインタです。

`PyInit_fibonacci`

PyMODINIT_FUNC PyInit_fibonacci(void){
    Py_Initialize();
    return PyModule_Create(&fibonacci_module_definition);
}

モジュールの初期化関数です。関数名は「PyInit_~」の形式である必要があります。

拡張モジュールをコンパイルする

できたコードをコンパイルするためにsetup.pyを使います。

setup.py

from setuptools import setup, Extension
setup(
        name = 'fibonacci',
        ext_modules = [
            Extension('fibonacci', ['fibonacci.c']),
            ]
        )

では、ビルドを行います。この作業はvenv等の仮想環境で行うことをお勧めします。

pip install -e .

使ってみる

速度比較のために次のコードを使います。

import time
import fibonacci    #C拡張

def fib(n):
    if n < 2:
        return 1;
    else:
        return fib(n-1) + fib(n-2)

if __name__ == '__main__':
    start = time.time()
    fib_py = [fib(i) for i in range(35)]
    elapsed_time = time.time() - start
    print(elapsed_time)

    start = time.time()
    fib = [fibonacci.fibonacci(i) for i in range(35)]
    elapsed_time = time.time() - start
    print(elapsed_time)

これを実行すると、Pythonの関数は6秒くらいなのに対してC拡張では0.08秒程度で処理が終わります。80倍くらいの性能向上が実現されています。

さらに最適化してみる

さらに実行速度を上げるために、拡張のコードをメモ化を使って書き換えます。変更箇所はfibonacci関数のみです。

long long dp[1000];
long long fibonacci(unsigned int n){
    if (n < 2){
        return 1;
    }else{
        if (dp[n-1] != 0 && dp[n-2] != 0){
            dp[n] = dp[n-1] + dp[n-2];
            return dp[n];
        }else{
            dp[n] = fibonacci(n-1) + fibonacci(n-2);
            return dp[n];
        }
    }
}

これでビルドしなおして先ほどのコードを実行してみると、4e-5秒くらいで終わります。