Posted at

RustでNumPyを拡張する

More than 1 year has passed since last update.


Pythonの利点


  • 対話的にデータの加工、解析、可視化 on Jupyter

  • グルー言語的用法


    • 公式・非公式によってPythonインターフェースはだいたい用意されている



  • pip/wheel、あるいはanacondaによる環境構築の容易さ


Pythonの問題点


  • Python自体の動作が低速



高速化の取り組み


  • Cython


    • コンパイルできるPythonに近いDSLを定義してコンパイルして実行する



  • Numba


    • PythonのコードをLLVMにコンパイルして実行する



  • Theano


    • Python上で埋め込みDSLによって計算を構築して、それを高速なバックエンドで実行する



  • FFI


    • 既に存在する高速に動作するライブラリ(C/Fortran)を呼び出す





Pythonの拡張とNumPyの拡張

2種類のC-APIの両方を使用する必要がある



どの言語でNumPyを拡張する?


  • C/C++


    • SWIG (NumPy C-APIはそのまま使う)



  • C++


    • Boost.Python (含Boost.NumPy)



  • Fortran


    • f2py



  • Rust


    • rust-cpython / rust-numpy





何故Rust?


  • Rustは現代的なシステム言語で数値計算に向いている(と思う)


    • Borrow Checkerによるデータ競合(data race)の排除

    • C++的なメモリモデル・move semantics (no GC)

    • 型クラス的なGenerics

    • 強力な型推論

    • Native Thread

    • LLVMによる最適化



  • Rustで書きたいから



rust-cpython

https://github.com/dgrunwald/rust-cpython


  • Python C-APIのRustラッパー


    • 参照カウントの管理

    • GILの抽象化

    • Pythonインタプリタの起動

    • Pythonへの関数の公開





RustからPythonの呼び出し

extern crate cpython;

use cpython::{Python, PyDict, PyResult};

fn main() {
let gil = Python::acquire_gil();
hello(gil.python()).unwrap();
}

fn hello(py: Python) -> PyResult<()> {
let sys = py.import("sys")?;
let version: String = sys.get(py, "version")?.extract(py)?;

let locals = PyDict::new(py);
locals.set_item(py, "os", py.import("os")?)?;
let user: String = py.eval("os.getenv('USER') or os.getenv('USERNAME')", None, Some(&locals))?.extract(py)?;

println!("Hello {}, I'm Python {}", user, version);
Ok(())
}



PythonからRustの関数の呼び出し

#[macro_use] extern crate cpython;

use cpython::{PyResult, Python};

// add bindings to the generated python module
// N.B: names: "librust2py" must be the name of the `.so` or `.pyd` file
py_module_initializer!(librust2py, initlibrust2py, PyInit_librust2py, |py, m| {
try!(m.add(py, "__doc__", "This module is implemented in Rust."));
try!(m.add(py, "sum_as_string", py_fn!(py, sum_as_string_py(a: i64, b:i64))));
Ok(())
});

// logic implemented as a normal rust function
fn sum_as_string(a:i64, b:i64) -> String {
format!("{}", a + b).to_string()
}

// rust-cpython aware function. All of our python interface could be
// declared in a separate module.
// Note that the py_fn!() macro automatically converts the arguments from
// Python objects to Rust values; and the Rust return value back into a Python object.
fn sum_as_string_py(_: Python, a:i64, b:i64) -> PyResult<String> {
let out = sum_as_string(a, b);
Ok(out)
}



rust-numpy

https://github.com/termoshtt/rust-numpy


  • GWに作った(/・ω・)/

  • rust-cpythonをベースにNumPy C-APIをRust側に公開


  • numpy.ndarrayに対応するPyArrayを導入する

  • rust-ndarrayのArrayに変換する



rust-ndarray

https://github.com/bluss/rust-ndarray


  • Rustにおける線形代数ライブラリ

  • NumPyと同様のstride方式のndarray

extern crate ndarray;

use ndarray::*;

// immutable example
fn axpy(a: f64, x: ArrayViewD<f64>, y: ArrayViewD<f64>) -> ArrayD<f64> {
a * &x + &y
}

// mutable example (no return)
fn mult(a: f64, mut x: ArrayViewMutD<f64>) {
x *= a;
}



rust-numpy

#[macro_use]

extern crate cpython;
extern crate numpy;

use numpy::*;
use cpython::{PyResult, Python, PyObject};

// wrapper of `axpy`
fn axpy_py(py: Python, a: f64, x: PyArray, y: PyArray) -> PyResult<PyArray> {
let np = PyArrayModule::import(py)?;
let x = x.as_array().into_pyresult(py, "x must be f64 array")?;
let y = y.as_array().into_pyresult(py, "y must be f64 array")?;
Ok(axpy(a, x, y).into_pyarray(py, &np))
}

// wrapper of `mult`
fn mult_py(py: Python, a: f64, x: PyArray) -> PyResult<PyObject> {
let x = x.as_array_mut().into_pyresult(py, "x must be f64 array")?;
mult(a, x);
Ok(py.None()) // Python function must returns
}

/* Define module "_rust_ext" */
py_module_initializer!(_rust_ext, init_rust_ext, PyInit__rust_ext, |py, m| {
m.add(py, "__doc__", "Rust extension for NumPy")?;
m.add(py, "axpy", py_fn!(py, axpy_py(a: f64, x: PyArray, y: PyArray)))?;
m.add(py, "mult", py_fn!(py, mult_py(a: f64, x: PyArray)))?;
Ok(())
});



setuptools-rust


  • rust-cpythonによる拡張をsetuptoolsで扱うためのライブラリ

setup.py

extensions/Cargo.toml
src/lib.rs
rust_ext/__init__.py


setup.py

from setuptools import setup

from setuptools_rust import RustExtension

setup(name='rust_ext',
version='1.0',
rust_extensions=[
RustExtension('rust_ext._rust_ext', 'extensions/Cargo.toml')],
packages=['rust_ext'],
zip_safe=False)



rust_ext/__init__.py

from ._rust_ext import *




rust-numpyの呼び出し

python setup.py install

import rust_ext

import numpy as np

x = np.array([1.0, 2.0])
y = np.array([2.0, 3.0])
rust_ext.axpy(3, x, y)

(*'▽')動いた!



最後に:Rustが辛いときは


  • rust-jp.slack.com

  • ja.stackoverflow.com/questions/tagged/rust

  • 神戸Rustの会