RustモジュールをPythonから実行する (PyO3)

Rust Python

最近少しずつRustに手を出し始めています。

CPUバウンドとなっている重い処理をPythonから切り出してコンバートするユースケースを考えたときに、PythonからRustを呼び出すための仕組みが必要となります。
今回はPythonからRustバイナリを実行する (あるいはその逆) ためのバインディングを行うPyO3について、エラトステネスのふるいを実装しながら紹介します。

  • 他の選択肢としてrust-cpythonもあるのですが、最近はPyO3の方がメジャーっぽいです
  • 実行環境はMac OS X (10.14) を想定しています

github.com

パッケージの作成

ライブラリテンプレートを使ってeratosthenesパッケージを作成します。
加えてPyO3はnightlyバージョンのRust環境が必要ですので、インストールして切り替えています。

$ cargo new eratosthenes --lib
$ cd eratosthenes

$ rustup install nightly 
$ rustup default nightly 
$ rustup toolchain list
stable-x86_64-apple-darwin
nightly-x86_64-apple-darwin (default)

Cargo.tomlにPyO3の依存関係を追記します。ここでは0.9.0-alpha.1を使っています。

...
[lib]
name = "eratosthenes"
crate-type = ["cdylib"]

[dependencies.pyo3]
version = "0.9.0-alpha.1"
features = ["extension-module"]

Rustの実装

src/lib.rsを編集し、エラトステネスのふるいを実装します。

  • 最初の2行はPyO3で使うモジュールが宣言されてます
    • pyo3::prelude以下にはこの後使うPyModuleやPyResultなどの基本モジュールが公開されています
  • #[pymodule]アトリビュート付きのeratosthenesがPythonモジュール (import eratosthenes) にバインディングされます
    • この後定義するget_prime_numbersをモジュールへ追加してます (eratosthenes.get_prime_numbers(10) で実行できます)
  • #[pyfunction]アトリビュート付きのget_prime_numbersがPython関数にバインディングされます
    • 返り値Vecは、Pythonではintのlistとなります
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;

#[pymodule]
fn eratosthenes(py: Python, m: &PyModule) -> PyResult<()> {
    m.add_wrapped(wrap_pyfunction!(get_prime_numbers))?;
    Ok(())
}

#[pyfunction]
fn get_prime_numbers(n: i32) -> PyResult<Vec<i32>> {
    let mut flags = Vec::new();
    for _ in 0..n+1 {
        flags.push(true);
    }

    let upper = (n as f32).sqrt().floor() as i32;
    for i in 2..upper+1 {
        if !flags[i as usize] {
            continue;
        }

        let prime = i;

        let mut j = prime * 2;
        while j <= n {
            flags[j as usize] = false;
            j += prime;
        }
    }

    let mut primes = Vec::new();
    for i in 2..n+1 {
        if flags[i as usize] {
            primes.push(i);
        }
    }

    Ok(primes)
}

ビルド

releaseオプション付きでビルドすると、ダイナミックライブラリliberatosthenes.dylibが出力されます。

$ cargo build --release
$ ls target/release
build/                 deps/                  examples/              incremental/           liberatosthenes.dylib

Pythonから実行

liberatosthenes.dylibを、実行するPythonファイルと同じディレクトリにeratosthenes.soとしてコピー&リネームまたはシンボリックリンクさせると、Pythonから実行できます。

import eratosthenes

primes = eratosthenes.get_prime_numbers(10)
print(primes)  # [2, 3, 5, 7]

速度の比較

試しにPythonでもエラトステネスのふるいを実装して速度をざっくり比較してみました。10000000の素数リストを10回計算してます。

import eratosthenes

def get_prime_numbers(n: int):
    flags = [True for _ in range(n+2)]

    upper = int(n ** 0.5)
    for i in range(2, upper+1):
        if not flags[i]:
            continue

        prime = i

        j = prime * 2
        while j <= n:
            flags[j] = False
            j += prime

    primes = []
    for i in range(2, n+1):
        if flags[i]:
            primes.append(i)

    return primes

if __name__  == "__main__":
    import sys
    use_rust = len(sys.argv) == 2 and sys.argv[1] == "--rust"
    n = 10000000
    
    for _ in range(10):
        if use_rust:
            primes = eratosthenes.get_prime_numbers(n)
        else:
            primes = get_prime_numbers(n)

結果は以下の通りで、Rustが約40倍速い結果となりました。実行可能なバイナリなので当然と言えばそれまでですが、やっぱりすごい。

$ time python src/eratosthenes.py
python src/eratosthenes.py  26.31s user 0.51s system 99% cpu 26.867 total

$ time python src/eratosthenes.py --rust
python src/eratosthenes.py --rust  0.64s user 0.17s system 96% cpu 0.834 total