用 Rust 扩展 Python | 数学建模加速尝试

Python 在数学建模的数值求解上可以说是砍瓜切菜，但是有一个小问题：部分场合有点慢，尤其是涉及列表、ndarray 来回倒腾，循环，隐函数求解，配合一个优化器之类的东西反复调用这个函数，运行起来奇慢无比。我们需要加速他，不然算不完啊！

Rust 在科学计算上并没有 Python 里大而全的 scipy, sympy, matplotlib 等库，生态不够发达。如果我们能够将二者优势结合起来，取 Python 的生态和 Rust 的速度，岂不美哉。

在开始之前需要强调，实际效果还是看重写部分的关键程度。只要不是 100% Rust，就应当评估原 Python 部分所占最终时间的比例。只有当时间比例超过 50% 时，改写才是有意义的。

本文经过测试，发现以下策略卓有成效：

(?) 对于大 ndarray 的 numpy 直接向量化操作（重写性能差不多，没必要）
(x200+) 长列表的复杂的数学计算和映射 ([f(i) for i in arr]) 用 Rust 重写
(x5-10) 热数学计算函数用 Rust 重写
(x8) 用 Rust 重写”优化函数“
(x4-10) 多个重任务的多线程分治

以下策略会适得其反：

(x0.3) 用 Python 的 multiprocessing.pool.Pool
(x0.9) 在 Rust 中频繁调用 Python 函数
(?) 少量轻任务的的多线程分治

测试全部在个人设备上进行 (Python 3.12.5, Windows10, AMD Ryzen 9 7950X)

构建系统

目前 Python 调用其他语言主要有这几个方向

cpython extension: 尽管仅在 CPython 系列上受支持，Python 科学计算生态几乎全部建立在其上。毕竟它直接支持 C call，支持自定义结构。
ctype: 广受支持，可以加载动态链接库，然而只能传递基本类型，一般对象就不行了，不少支持的类型比如函数、数组等都要经过很重的兼容性重建，不适合作为热函数
cffi/uniffi: 太新了，不是很稳定
rpc: 在异步环境中还是很不错的，阻塞的计算密集型科学计算就不太适合了

我们选基于 native extension 的 PyO3。PyO3 (doc) 是一个相当完善的 Python-Rust binding。它提供了无痛的模块编写（写过 CPython 模块的都知道手搓有多痛…）

PyO3 推荐使用脚手架工具 maturin ~~(マチュリン)~~ 。它支持 pyo3, cffi 和 uniffi 三种模式，建议还是用 pyo3 模式。但是我们并不需要发布一个 wheel，只要在自己和朋友电脑的 python 里能 import 就行。完全可以不需要脚手架。

CPython 的 import 除了能引入一个 py 文件或者含有 __init__.py 的文件夹，也可以引入一个动态链接库。不过需要注意运行 Python 版本要与编译的一致，否则会报 “ImportError: DLL load failed while importing XXX”

只要把编译出来的动态链接库重命名为 <模块名>.pyd (Windows) 或者 <模块名>.so (Linux) 就可以了。

cargo new --lib <模块名>
cargo add [email protected] # 注意版本保持一致
cargo add [email protected] -F nalgebra

lib.rs

use pyo3::prelude::*;

#[pyfunction]
fn sum(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

#[pymodule]
fn t(m: &Bound<'_, PyModule>) -> PyResult<()> {
    macro_rules! add_functions {
        ($($f:ident),*) => {$(
            m.add_function(wrap_pyfunction!($f, m)?)?;
        )*};
    }
    add_functions!(sum);
    Ok(())
}

Cargo.toml

[lib]
name = "p"
crate-type = ["cdylib"]

然后把编译产物动态链接库重命名为 t.pyd/so（与模块名同名即可），在该目录下运行 Python

import t
t.sum(1,2)=='3'

其实你可以定义多个 pymodule，然后软/硬链接它到对于模块名的文件上，就可以单文件实现多个模块了。

经过测试，一个自动生成的 pyfunction 和 Python 内函数调用开销相同（15ns/op）。超快！

如果你要分享给你朋友，而且你们用的不是一个版本，只要把他版本的 Python 装在随便一个目录里（一个框都不勾），找到 libs 目录的绝对路径填到下面。如果是 linux 用户可以用一分钟编译一份 Python （注意 enable-shared），源代码目录作为 PYO3_CROSS_LIB_DIR （可能还要加个 RUSTFLAGS=’-L${DIR}’），不需要安装。这主要是为了确定版本并链接对应版本的 Python 动态库。

PYO3_CROSS_PYTHON_VERSION=3.8 PYO3_CROSS_LIB_DIR="C:\Program Files\Python38\libs" cargo build -r

numpy

PyO3 提供了 numpy 支持，可以零拷贝转换为 nalgebra::Matrix 或是 ndarray::ArrayBase 从而在 Rust 内计算。

use numpy::PyReadonlyArray1;
#[pyfunction]
fn n(arr: PyReadonlyArray1<f64>)->PyResult<f64> {
    let mut sum=0.;
    for i in arr.as_slice()? {
        sum+=*i;
    }
    Ok(sum)
}

十分轻松写意。

如果要新建对象，需要用到 python 运行时。Python 对象生命周期完全由 Python GC 管理，因此需要从 Python 运行时创建、销毁对象。

PyO3 的 pyfunction 支持首个参数为 Python 的函数，这样后面的参数会被依次解析，仿佛只有后面的参数一样。这样就可以通过这个 Python 对象创建销毁对象了。btw PyO3 为了销毁安全性，引入了 Bound 来为每个对象指针绑定一个 Python 指针，这样就可以自动销毁，不用手动调用从 Python 运行时销毁了。因此这个 Python 参数仅当需要创建新对象时才应使用。

内置函数

use numpy::{PyArrayDyn, PyArrayMethods};
use pyo3::types::PyTuple;
#[pyfunction]
fn twos<'py>(py: Python<'py>, shape: Bound<'py, PyTuple>)->PyResult<Bound<'py, PyArrayDyn<f64>>> {
    let shape = shape.as_slice();
    let mut s = Vec::with_capacity(shape.len());
    for i in shape {
        s.push(i.extract::<usize>()?)
    }
    let mut len=1;
    s.iter().for_each(|i| len*=i);
    let arr = unsafe{PyArrayDyn::new_bound(py, s.as_slice(), false)};
    let raw = unsafe{core::slice::from_raw_parts_mut(arr.data(), len)};
    raw.fill(2.);
    Ok(arr)
}

经过 benchmark，小形状下比 np.ones 快不少，大形状慢点，可能是 slice::fill 还不够快吧。

简单数学计算

bench_total=100

import numpy as np
arr = np.arange(2,100086)

def my_map(x):
    if x<=0:
        x=-x
    while x>2:
        x/=2
    return x
def run_map():
    return np.array([my_map(i) for i in arr])

def my_map(x):
    y=x
    x=y
    if x<=0:
        x=-x
    while x>2:
        x/=2
    return x

from t import map_rs, map_rspy, my_map_rs

def run_map_rs():
    return np.array([my_map_rs(i) for i in arr])
   
if __name__ == '__main__':
	bench("python", run_map)
	bench("rewrite-map", run_map_rs)
	bench("rewrite-list", lambda: map_rspy(my_map, arr))
	bench("rust", lambda: map_rs(arr))

use pyo3::prelude::*;
use numpy::{PyArray1, PyArrayMethods};

#[pyfunction]
fn map_rs<'a>(py: Python<'a>, arr: &Bound<'a, PyArray1<i64>>)->PyResult<Bound<'a, PyArray1<f64>>> {
    let dims = arr.dims();
    let (new_arr, src, dst);
    unsafe {
        new_arr = PyArray1::new_bound(py, dims, false);
        src = arr.as_slice()?;
        dst = new_arr.as_slice_mut()?;
    }
    for (i,j) in src.iter().zip(dst.iter_mut()) {
        *j=my_map_rs(*i);
    }
    Ok(new_arr)
}

#[pyfunction]
fn my_map_rs(i:i64)->f64 {
    let mut i = i as f64;
    if i<0. {i=-i;}
    while i>2. {i/=2.}
    i
}

#[pyfunction]
fn map_rspy<'a>(py: Python<'a>, f:&Bound<'a, PyAny>, arr: &Bound<'a, PyArray1<i64>>)->PyResult<Bound<'a, PyArray1<f64>>> {
    let dims = arr.dims();
    let (new_arr, src, dst);
    let my_map_py = move |i:&i64|->PyResult<f64> {
        Ok(f.call1((*i,))?.extract::<f64>()?)
    };
    unsafe {
        new_arr = PyArray1::new_bound(py, dims, false);
        src = arr.as_slice()?;
        dst = new_arr.as_slice_mut()?;
    }
    for (i,j) in src.iter().zip(dst.iter_mut()) {
        *j=my_map_py(i)?;
    }
    Ok(new_arr)
}

#[pymodule]
fn t(m: &Bound<PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction_bound!(map_rs, m)?)?;
    m.add_function(wrap_pyfunction_bound!(map_rspy, m)?)?;
    m.add_function(wrap_pyfunction_bound!(my_map_rs, m)?)?;
    Ok(())
}

# arr = np.arange(2,186)
python: 77.608us/op
rewrite-map: 15.700us/op
rewrite-list: 43.753us/op
rust: 0.500us/op

# arr = np.arange(2,100086)
python: 89.391ms/op
rewrite-map: 9.463ms/op
rewrite-list: 56.068ms/op
rust: 0.340ms/op

可见重写数学函数能有 5-10 倍加速，重写列表映射函数能有 1.5-2 倍加速，而减少跨语言调用、两部分同时重写后加速了上百倍。这是一个好的开始！

如果最终拷贝回 ndarray (slice::copy_from_slice) 的开销和失去 blas 的性能损失可以接受的话，用 nalgebra 和 ndarray 还是好写又快的。毕竟静态类型加编译，可能还有编译器的 simd，性能损失不会超过 1.5 倍的。

argmin

argmin 是一个纯 Rust 实现的数值优化库，提供了与 rust 线性代数库 nalgebra 与张量库 ndarray 兼容的接口。

它提供了 Solver Trait 的各种算法，和一个 Executer Struct 执行器，需要我们自定义 Problem。

每个算法会依赖一些 Trait，比如雅可比矩阵，原函数，梯度，黑塞矩阵等等。我们的 Problem 只要实现了对应 Trait 集合就可以调用一个算法。

例如一元函数，我们使用 BrentOpt 算法。

use argmin::core::{CostFunction, Executor, Gradient, Hessian};
use argmin::solver::brent::BrentOpt;
struct OneOne<F> (F);

impl<F: Fn(f64)->f64> CostFunction for OneOne<F> {
    type Param = f64;
    type Output = f64;
    fn cost(&self, param: &Self::Param) -> Result<Self::Output, argmin::core::Error> {
        Ok((self.0)(*param))
    }
}

fn main() {
    let problem = OneOne(|x:f64| x*x);
    let solver = BrentOpt::new(-10., 10.);
    let runner = Executor::new(problem, solver);
    println!("{}", runner.configure(|c|c.max_iters(1000)).run().unwrap());
}

会输出

OptimizationResult:
    Solver:        BrentOpt
    param (best):  0.0
    cost (best):   0
    iters (best):  2
    iters (total): 6
    termination:   Solver converged
    time:          6.1µs

有点慢，也不算太慢。迭代次数够少，适合单次函数运行较为缓慢的情况。

有朋友可能会想，要不把 python 实现的算法函数传入 pyo3，这样不用引入 rust 的算法库；然而实测这样会慢 15%。

一个简单的测试代码如下

use pyo3::prelude::*;
use argmin::{core::{CostFunction, Executor}, solver::brent::BrentRoot};

#[pyfunction]
fn solve<'a>(py: Python<'a>, fsolve: &Bound<'a, PyAny>)->PyResult<Bound<'a, PyAny>> {
    #[pyfunction]
    fn to_solve(x:f64)->f64 {x*x-1.}
    let tosolve_py = wrap_pyfunction_bound!(to_solve, py)?;
    fsolve.call((tosolve_py, 10086.), None)?.get_item(0)
}

#[pyfunction]
fn solve_argmin()->f64 {
    struct P;
    impl CostFunction for P {
        type Param = f64;
        type Output = f64;
        fn cost(&self, param: &Self::Param) -> Result<Self::Output, argmin::core::Error> {
            Ok(param*param-1.)
        }
    }
    let solver = BrentRoot::new(0., 10086., 1e-10);
    let runner = Executor::new(P, solver);
    let r = runner.run().unwrap();
    r.state.best_param.unwrap()
}

#[pymodule]
fn t(m: &Bound<PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction_bound!(solve, m)?)?;
    m.add_function(wrap_pyfunction_bound!(solve_argmin, m)?)?;
    Ok(())
}

from scipy.optimize import fsolve
from t import solve,solve_argmin
import time

total=10000

def bench(name, fn):
    start_time = time.time_ns()
    for _ in range(total):
        fn()
    stop_time = time.time_ns()
    print("{}: {:.3f}us/op; result: {:.8f}".format(
        name, (stop_time-start_time)/total/1000, fn()
    ))

to_solve=lambda x:x*x-1

def solve_py(solver):
	return fsolve(to_solve, 10086.)[0]

bench("solve", lambda: solve(fsolve))
bench("solve_argmin", lambda: solve_argmin())
bench("solve_py", lambda: solve_py(fsolve))

输出

solve: 46.392us/op; result: 1.00000000
solve_argmin: 5.009us/op; result: 1.00000000
solve_py: 40.130us/op; result: 1.00000000

也就 8 倍速度，而且有误差，需要确定解边界，且需要满足 $f(left)f(right)<0$。

真实的多解情况中，可以尝试先线性搜索可能的区间，然后每个可能区间跑一次。

多线程

不只是 Rust，Python 里也可以用这个技巧加速计算。只要是无关的任务列表，比如复杂的列表逐元素计算，多个实例计算，都可以用多线程加速计算。

但是，在继续之前需要考虑一下，是否值得使用多线程。一次线程的建立就要 100ns 以上，线程间任务分配和同步都较为耗时。如果预估总的计算量在 5us 以下或者单次计算在 10ns 以下，用多线程吃力不讨好。

Python 里自带一个 multiprocessing.pool 可以并行化计算 list。

我们沿用上面 numpy 数学计算的例子

from multiprocessing.pool import Pool
def run_map_multi():
    pool = Pool(processes=16)
    return np.array(pool.map(my_map, arr))

if __name__ == '__main__':
	bench("python-par", run_map_multi)

结果：

# arr = np.arange(2,100086)
python: 91.279ms/op
python-par: 250.499ms/op

多线程被单线程打爆了！这是由于

Python 本身由于 GIL 的设计只能单线程，因此其通过多进程实现。这就开销更大了
my_map 太快了，本身只要 0.5 us 就可以计算一次
线程间内容同步很慢

当我们把 my_map 手动放慢 100 倍，可以看到

python: 649.737ms/op
python-par: 185.958ms/op

确实是多线程快了。

当然，如果 Rust 里多线程呢（基于 rayon）

#[pyfunction]
fn map_rs_par<'a>(py: Python<'a>, arr: &Bound<'a, PyArray1<i64>>)->PyResult<Bound<'a, PyArray1<f64>>> {
    use rayon::prelude::*;
    let dims = arr.dims();
    let (new_arr, src, dst);
    unsafe {
        new_arr = PyArray1::new_bound(py, dims, false);
        src = arr.as_slice()?;
        dst = new_arr.as_slice_mut()?;
    }
    dst.into_par_iter().zip(src).for_each(|(d,s)|{
        *d = my_map_rs(*s);
    });
    Ok(new_arr)
}

plain: 321.209us/op
rayon: 125.000us/op

两倍多速度。

当然，如果把计算任务加重，从单次 34ns 到 184ns：

fn my_map_rs(i:i64)->f64 {
    let mut i = i as f64;
    if i<0. {i=-i;}
    while i>2. {i=(i+i.sin())*0.5}
    i
}

效果就更显著了：（可见，还是计算越繁重，用多线程的优势越大。）

plain: 15.129ms/op
rayon: 1.400ms/op

即使是 10086 个数据也是多线程速度在单线程 8 倍左右。在 1e3 以下的量级优势就不明显了。

最终的代码

use pyo3::prelude::*;
use numpy::{PyArray1, PyArrayMethods};

#[pyfunction]
fn my_map_rs(i:i64)->f64 {
    let mut i = (i as f64).abs();
    while i>2. {i=i/2.}
    i
}

fn map_rs_with<'a>(py: Python<'a>, arr: &Bound<'a, PyArray1<i64>>, f: impl FnOnce(&mut [f64], &[i64]))->PyResult<Bound<'a, PyArray1<f64>>> {
    let dims = arr.dims();
    let (new_arr, src, dst);
    unsafe {
        new_arr = PyArray1::new_bound(py, dims, false);
        src = arr.as_slice()?;
        dst = new_arr.as_slice_mut()?;
    }
    f(dst, src);
    Ok(new_arr)
}

#[pyfunction]
fn map_rs<'a>(py: Python<'a>, arr: &Bound<'a, PyArray1<i64>>)->PyResult<Bound<'a, PyArray1<f64>>> {
    map_rs_with(py, arr, |dst, src| {
        for (i,j) in src.iter().zip(dst.iter_mut()) {
            *j=my_map_rs(*i);
        }
    })
}

#[pyfunction]
fn map_rs_par<'a>(py: Python<'a>, arr: &Bound<'a, PyArray1<i64>>)->PyResult<Bound<'a, PyArray1<f64>>> {
    map_rs_with(py, arr, |dst, src| {
        use rayon::prelude::*;
        dst.into_par_iter().zip(src).for_each(|(d,s)|{
            *d = my_map_rs(*s);
        });
    })
}

#[pymodule]
fn t(m: &Bound<PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction_bound!(map_rs_par, m)?)?;
    m.add_function(wrap_pyfunction_bound!(map_rs, m)?)?;
    Ok(())
}

后记

本次探索过程还是十分有趣的。PyO3 的上手惊人地舒适，依赖也很少，写起 Python 模块来还是蛮顺手的。

在数学建模场景下，时间有限，如果确定了数值模拟类题目，最好从开始就把 cpu 密集部分用 Rust 写，以免重构时间过多。