PyO3 是 Python 的 Rust 绑定。

虚拟环境与 maturin

创建一个 Python 虚拟环境（更好的管理）

1
2
mkdir simple_trig
pipenv install

激活虚拟环境，安装 maturin

1
2
pipenv shell
pipenv install maturin

初始化 maturin

1
maturin init

会自动生成 Rust 的项目，生成的 lib.rs 如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
use pyo3::prelude::*;

/// Formats the sum of two numbers as string.
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

/// A Python module implemented in Rust.
#[pymodule]
fn simple_trig(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

只需要运行 maturin develop，就可以自动打包和安装到 Python 的虚拟环境。

1
2
3
4
% python
>>> import simple_trig
>>> simple_trig.sum_as_string(5, 20)
'25'

下面仿照 C++ 的例子，见 使用 C++ 扩展 Python：调用 C++ DLL，复刻成 Rust。

新建 Python 脚本

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from random import random
from time import perf_counter

COUNT = 500000  # Change this value depending on the speed of your computer
DATA = [(random() - 0.5) * 3 for _ in range(COUNT)]

e = 2.7182818284590452353602874713527

def sinh(x):
    return (1 - (e ** (-2 * x))) / (2 * (e ** -x))

def cosh(x):
    return (1 + (e ** (-2 * x))) / (2 * (e ** -x))

def tanh(x):
    tanh_x = sinh(x) / cosh(x)
    return tanh_x

def test(fn, name):
    start = perf_counter()
    result = fn(DATA)
    duration = perf_counter() - start
    print('{} took {:.3f} seconds\n'.format(name, duration))

    for d in result:
        assert -1 <= d <= 1, " incorrect values"

if __name__ == "__main__":
    print('Running benchmarks with COUNT = {}'.format(COUNT))

    test(lambda d: [tanh(x) for x in d], '[tanh(x) for x in d] (Python implementation)')

Rust 实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
use pyo3::prelude::*;

/// Formats the sum of two numbers as string.
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

const E: f64 = 2.7182818284590452353602874713527;

fn sinh_impl(x: f64) -> f64 {
    (1.0 - f64::powf(E, -2.0 * x)) / (2.0 * f64::powf(E, -x))
}

fn cosh_impl(x: f64) -> f64 {
    (1.0 + f64::powf(E, -2.0 * x)) / (2.0 * f64::powf(E, -x))
}

#[pyfunction]
fn fast_tanh(x: f64) -> PyResult<f64> {
    Ok(sinh_impl(x) / cosh_impl(x))
}

/// A Python module implemented in Rust.
#[pymodule]
fn simple_trig(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    m.add_function(wrap_pyfunction!(fast_tanh, m)?)?;
    Ok(())
}

加入测试函数

1
2
3
4
5
6
7
from simple_trig import fast_tanh

if __name__ == "__main__":
    print('Running benchmarks with COUNT = {}'.format(COUNT))

    test(lambda d: [tanh(x) for x in d], '[tanh(x) for x in d] (Python implementation)')
    test(lambda d: [fast_tanh(x) for x in d], '[fast_tanh(x) for x in d] (Rust extension)')

令我震惊的是，Rust 居然更慢……

1
2
3
4
Running benchmarks with COUNT = 500000
[tanh(x) for x in d] (Python implementation) took 0.234 seconds

[fast_tanh(x) for x in d] (Rust extension) took 0.282 seconds

好吧，一定是哪里出问题了。

（没准，这就是事实？）

小结

pybind11 与 pyo3

Python 3.11 下，PyO3 居然会更慢，但使用 Pybind11，确实是更快了。我也不知道是不是我哪里没有开优化。