用Rust帮Python加加速

背景

长期以来,Python由于易上手,有GC且生态强大等特点被广泛使用,可是渐渐的人们也发现了它的不足,解释型语言的运行速度终究比不过编译型,况且由于Python设计时的动态数据类型一切皆对象(内存都分配在堆上)等思想,也导致了运行速度缓慢.

随着实时性要求的不断提升,在一些计算量大要求快速响应的场景传统的Python就很难满足要求,所以随之慢慢有了各种解决办法:

用更高效的解释器
用jit即时编译加速
改写成Cython加速
对GIL动手,提高多线程性能
…

其中目前使用最广泛,最有效的应该是jit与Cython这两种方案,jit即时编译可以将部分需要解释的代码直接转为机器码从而实现加速(减少解释时开销);而Cython更绝直接将Python原地升级,得到一个Cython这个Python与C的混血,可以通过Cython将代码翻译成C/C++的代码再编译成动态库文件供使用.可是这样就会造成Python原本的语法被改的“四不像”,这些后面慢慢再谈.

既然都用上动态库了,为什么不直接使用C++或者其他高效语言实现,然后供Python调用呢?说到底,Cython不也是翻译成C/C++的代码编译使用,只是为了方便Python开发人员才设计了这种类似于Python的语法.如果熟悉其他语言的话完全可以直接使用其他语言实现而不影响Python的基本语法.所以今天就来讨论一下关于使用Rust对Python计算进行加速的问题.

Rust加速Python计算

首先,来看看Rust实现和Python实现基本的速度对比,目标是求斐波那契数列第n项的值.其中实现均采用递归调用,为了突出时间差异这里求第30项的值并重复50次

Python的实现与耗时如下:

import time

def fib(n:int) ->int:
    assert n>=0
    if n <= 1:
        return n
    return fib(n-1)+fib(n-2)

def main(test_times=50):
    start = time.time()
    for _ in range(test_times):
        fib(30)
    print(f"time cost {time.time()-start} s")

if __name__ == '__main__':
    main()

Rust的实现与耗时如下:

use std::time;

fn fib(n:i32)->u64{
    if n<=0{
        panic!("{} must be a postive number!",n);
    }
    match n{
        1|2 => 1,
        _ => fib(n-1) + fib(n-2)
    }
}

fn main() {
    let test_times = 50;
    let start = time::Instant::now();
    for i in 0..test_times{
        fib(30);
    }
    println!("time cost {:?}",start.elapsed())
}

这差异,足足一百多倍.那看来使用Rust提速是完全可行的,那怎么将Rust与Python相结合呢?或者如何把Rust的代码编译供Python调用,这个时候可以使用pyo3,首先安装一下maturin工具pip install maturin,然后配置一下项目的Cargo.toml

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

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

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

[dependencies]
pyo3 = { version = "0.19.2", features = ["extension-module"] }

这里编译类型就设置为lib,关于多种不同lib类型的区别可以去看看Rust专栏之前的内容.这里的name就是未来Python中调用的名字,下面再编写lib.rs

use pyo3::prelude::*;
use pyo3::wrap_pyfunction;

#[pyfunction]
pub fn fib(n:i32)->u64{
    if n<=0{
        panic!("{} must be a postive number!",n);
    }
    match n{
        1|2 => 1,
        _ => fib(n-1) + fib(n-2)
    }
}

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

逻辑代码基本没有更改,只是添加了Rust实现Python module的代码,这里的module name必须和toml中设置的name保持一致,否则也会无法导入.

最后运行maturin develop就可以实现编译,给Python调用了.

加速对比

现在,我们已经实现了Rust的加速,是不是非常简单而且调用的时候可以使用原本的Python语法而不用进行任何更改.下面就来对比一下原始,numba,Cython,Rust四种方式的速度对比.

其中Cython实现cpy_fib.pyx如下

cpdef int c_fib(n:int):
    assert n>0
    if n in [1, 2]:
        return 1
    else:
        return c_fib(n - 1) + c_fib(n - 2)

然后写setup进行编译

from distutils.core import setup,Extension
from Cython.Build import cythonize

setup(
    ext_modules=cythonize(Extension(
        'cpy_fib',
        sources=['./cpy_fib.pyx'],
        language='c'
    )),
)

运行python setup.py build_ext --inplace编译,最后整体对比

import speed_python
import time
from cpy_fib import c_fib
from numba import jit


@jit(nopython=True)
def fib_jit(n: int) -> int:
    assert n > 0
    if n in [1, 2]:
        return 1
    else:
        return fib_jit(n - 1) + fib_jit(n - 2)

def fib(n: int) -> int:
    assert n > 0
    if n in [1, 2]:
        return 1
    else:
        return fib(n - 1) + fib(n - 2)


def test_speed(func,func_name:str,test_times=50):
    start = time.time()
    for _ in range(test_times):
        func(30)
    print(f"{func_name} speed up time cost {time.time() - start} s")


def main(test_times=50):
    test_speed(fib,"origin python")
    test_speed(fib_jit,"numba python")
    test_speed(speed_python.fib,"rust")
    test_speed(c_fib,"Cython")


if __name__ == '__main__':
    main()

当然如果希望进一步加速,我们还是有办法的.使用maturin develop --release生成Rust的调用,再来比较一下计算耗时,速度又提升了一倍多