用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的调用,再来比较一下计算耗时,速度又提升了一倍多

加速方法	耗时(ms)
原始	4144
Numba	441
Cython	804
Rust	178 / 61

这加速对比也证明了在一些计算任务中Rust能更高效的实现,并且不影响原始Python的代码语法或者结构,只需要编译调用.完全可以分配不同开发人员同时开发,最后整合测试调用即可.

如果你对Python感兴趣，想要学习python，这里给大家分享一份Python全套学习资料，都是我自己学习时整理的，希望可以帮到你，一起加油！

😝有需要的小伙伴，可以V扫描下方二维码免费领取🆓

​

1️⃣零基础入门

① 学习路线

对于从来没有接触过Python的同学，我们帮你准备了详细的学习成长路线图。可以说是最科学最系统的学习路线，你可以按照上面的知识点去找对应的学习资源，保证自己学得较为全面。

② 路线对应学习视频

还有很多适合0基础入门的学习视频，有了这些视频，轻轻松松上手Python~

③练习题

每节视频课后，都有对应的练习题哦，可以检验学习成果哈哈！

2️⃣国内外Python书籍、文档

① 文档和书籍资料

3️⃣Python工具包+项目源码合集

①Python工具包

学习Python常用的开发软件都在这里了！每个都有详细的安装教程，保证你可以安装成功哦！

②Python实战案例

光学理论是没用的，要学会跟着一起敲代码，动手实操，才能将自己的所学运用到实际当中去，这时候可以搞点实战案例来学习。100+实战案例源码等你来拿！

③Python小游戏源码

如果觉得上面的实战案例有点枯燥，可以试试自己用Python编写小游戏，让你的学习过程中增添一点趣味！

4️⃣Python面试题

我们学会了Python之后，有了技能就可以出去找工作啦！下面这些面试题是都来自阿里、腾讯、字节等一线互联网大厂，并且有阿里大佬给出了权威的解答，刷完这一套面试资料相信大家都能找到满意的工作。