使用pyo3 加速pypy和cpython¶

pyo3¶

https://github.com/PyO3/pyo3

maturin¶

https://github.com/PyO3/maturin

前置¶

版本¶

cpython版本: 3.8.6

pypy版本: 3.7-v7.3.5

rust版本: 1.57.0-nightly

工作目录设置¶

cd ~/Documents
mkdir pypy_cpython_pyo3     # 创建工作目录
cd pypy_cpython_pyo3
cargo new --lib BubbleSort  # 创建rust 项目
touch test.py               # 创建py 文件
├── BubbleSort
│   ├── Cargo.lock
│   ├── Cargo.toml
│   └── src
│       └── lib.rs
└── test.py

rust工程代码¶

修改 Cargo.toml¶

[package]
name = "BubbleSort"
version = "0.1.0"
edition = "2021"

[lib]
name = "BubbleSort"
crate-type = ["cdylib"]     # 注意这里一定要有 "cdylib"

[dependencies.pyo3]
version = "0.14.5"
features = ["extension-module"]

编写我们的主要功能¶

具体pyo3 的使用方法可以查阅文档, 这里只贴代码

use pyo3::prelude::*;


#[pyfunction]
fn bubble(mut li: Vec<u32>) -> PyResult<Vec<u32>> {
    let n = li.len();
    for i in 0..n{
        for j in 0..(n-i-1) {
            if li[j] > li[j+1]{
                let temp = li[j + 1];
                li[j + 1] = li[j];
                li[j] = temp;
            }
        }
    }
    Ok(li)
}


#[pymodule]
// 这个函数的名字必须与`Cargo.toml`中的`lib.name`设置相匹配，否则Python将无法导入该模块。
fn BubbleSort(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(bubble, m)?)?;
    Ok(())
}

python代码¶

# test.py


import random
import time
from BubbleSort import *


def get_li(num_count):
    return [random.randint(1, 100) for _ in range(num_count)]


def run(for_count, num_count):
    for i in range(for_count):
        li = get_li(num_count)
        bubble(li)


if __name__ == '__main__':
    print(dir())

    start = time.perf_counter()
    run(10, 1_0000)
    print(round(time.perf_counter() - start, 5))

手动编译¶

cpython + pyo3¶

$ cd BubbleSort               # 进入 rust工作目录
$ cargo build --release       # release 编译

# 移动编译后的库文件到指定位置 (!注意文件名变化)
$ mv target/release/libBubbleSort.dylib ../BubbleSort.so

# 构建完毕 执行
$ cd ..                       # 返回工作目录
$ python3 test.py             # 运行 py文件
>>> ['__annotations__', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__spec__', 'bubble', 'get_li', 'random', 'run', 'time']
>>> 0.67992

正确执行之后, 请删除 rust工程中的target目录, 以及工作目录中的BubbleSort.so 文件, 还原一个干净的工作目录

使用 maturin 自动构建¶

cpython + pyo3¶

这里很简单, 按照官方github上的示例即可(官方使用的是虚拟环境, 当然也可以不使用)

$ cd BubbleSort               # 进入 rust工程目录
$ python3 -m venv .env        # 创建虚拟环境
$ source .env/bin/activate    # 激活虚拟环境
$ pip install --upgrade pip   # 升级 pip
$ pip install maturin         # 安装 maturin

# 开始自动构建
$ maturin develop --release   # 使用以 develop 构建(会在.env/site-packages自动生成包)

# 构建完毕 执行
$ cd ..                       # 返回工作目录
$ python3 test.py             # 运行 py文件
>>> ['BubbleSort', '__annotations__', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__spec__', 'bubble', 'get_li', 'random', 'run', 'time']
>>> 0.67553

正确执行之后, 请删除 rust工程中的target目录, 删除 rust工程中的 .env 文件夹还原一个干净的工作目录, 并退出此虚拟环境

pypy + pyo3¶

中间也踩过不少坑(因为多删了一行Cargo.toml中的代码, 导致浪费了半天功夫), 一开始我使用的手动构建, 且设置了 PYO3_PYTHON和PYTHON_SYS_EXECUTABLE, 但是打包后的so, pypy解释器一直找不到内部的符号, 后面看着git上面pyo3构建cpython动态库的示例, 使用的是maturin这个工具构建的, 我寻思这 pypy 也试试用这个工具试试会怎样? 结果成功了, 蛋疼

注意, 我这里使用pypy3.6的版本在创建虚拟环境的时候创建失败, 换成pypy3.7的版本, 一切正常

操作和cpython差不多, 只不过创建虚拟环境的时候, 注意一下

$ cd BubbleSort               # 进入 rust工程目录
$ pypy_path/pypy3 -m venv .env# 创建虚拟环境
$ source .env/bin/activate    # 激活虚拟环境
$ export PYTHONPATH=.env/site-packages/
$ pip install --upgrade pip   # 升级 pip
$ pip install maturin         # 安装 maturin

# 开始自动构建
$ maturin develop --release   # 使用以 develop 构建(会在.env/site-packages自动生成包)

# 构建完毕 执行
$ cd ..                       # 返回工作目录
$ python3 test.py             # 运行 py文件

>>> ['BubbleSort', '__annotations__', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__spec__', 'bubble', 'get_li', 'random', 'run', 'time']
>>> 0.66007

maturin 的基本使用¶

maturin list-python: 查看哪些环境被选中

maturin develop: 根据env环境构建crate, 并将其作为一个python模块直接安装在当前的virtualenv中的site-packages中

maturin build: 构建一个 whl 的文件在 target/wheels 文件夹中, 可以使用 pip install /target/wheels/xxx.whl 安装 (pypy环境下使用该功能, 有个bug, 生成的whl文件还是cpython使用的, pypy是用不了, 使用 -i指定pypy解释器, 编译出错, 设置环境变量但是无用)