RustでPythonよりもパフォーマンスが出ない

Question

Rustを使った数値計算のプログラムを書いています。その中で、昔書いた Pythonで使う用のコードをRust側で再利用しているのですが、そこでのパフォーマンスがかなり低く、解決方法を探しています。

環境

• Rust 1.55.0
• WSL2 (Windows 10)

詳細

Python用に書いたコードは以下のようになっています。

lib.rs

mod kk;

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

use kk::kk::{real2imag_helper, imag2real_helper, kk_transform};

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

#[pyfunction]
fn real2imag(x: Vec<f64>, y: Vec<f64>) -> PyResult<Vec<f64>> {
    kk_transform(x, y, real2imag_helper)
}

/// Calculate the real part from the imaginary part
///
/// * `x` - independent variable (ex. energy or frequency)
/// * `y` - dependent variable (ex. conductivity or permittivity)
#[pyfunction]
fn imag2real(x: Vec<f64>, y: Vec<f64>) -> PyResult<Vec<f64>> {
    kk_transform(x, y, imag2real_helper)
}

kk.rs

use std::sync::{Mutex, Arc};
use std::thread;
use std::f64::consts::PI;

use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;

pub fn kk_transform<F>(x: Vec<f64>, y: Vec<f64>, f: F) -> PyResult<Vec<f64>>
    where F: Fn(&Vec<f64>, &Vec<f64>, usize) -> f64,
          F: Send + Copy + 'static
{
    let thread_num = 16;
    let mut handles: Vec<thread::JoinHandle<()>> = Vec::new();
    let mut result: Arc<Vec<Mutex<f64>>> = Arc::new(
        vec![0.0; y.len()]
            .into_iter()
            .map(|x| Mutex::new(x))
            .collect()
    );

    for i in 0..thread_num {
        let x = x.clone();
        let y = y.clone();

        let result = Arc::clone(&mut result);
        let handle = thread::spawn(move || {
            for j in x.len()*i/thread_num..x.len()*(i+1)/thread_num {
                let mut val = result[j].lock().unwrap();
                *val = 2.0 / PI * f(&x, &y, j);
            }
        });

        handles.push(handle);
    }

    for handle in handles {
        handle.join().unwrap();
    }

    let result = Arc::try_unwrap(result).unwrap();

    let mut output = vec![0.0; y.len()];
    for (i, val) in result.iter().enumerate() {
        output[i] = *val.lock().unwrap();
    }

    Ok(output)
}

#[allow(dead_code)]
pub fn real2imag_helper(x: &Vec<f64>, y: &Vec<f64>, num: usize) -> f64 {
    let mut result = 0.0;
    let diff = x[1] - x[0];

    let base = x[num];

    for (xx, yy) in x.iter().zip(y.iter()) {
        if *xx == base {
            continue;
        }
        result -= base * yy / (xx * xx - base * base) * diff;
    }

    result
}

#[allow(dead_code)]
pub fn imag2real_helper(x: &Vec<f64>, y: &Vec<f64>, num: usize) -> f64 {
    let mut result = 0.0;
    let diff = x[1] - x[0];

    let base = x[num];

    for (xx, yy) in x.iter().zip(y.iter()) {
        if *xx == x[num] {
            continue;
        }
        result += xx * yy / (xx * xx - base * base) * diff;
    }

    result
}

これをRustで再利用するために、コピペして今書いているプロジェクトのコードに導入しました。

use std::f64::consts::PI;
use std::sync::{Arc, Mutex};
use std::thread;

pub fn imag2real(x: Vec<f64>, y: Vec<f64>) -> Vec<f64> {
    kk_transform(x, y, imag2real_helper)
}

fn kk_transform<F>(x: Vec<f64>, y: Vec<f64>, f: F) -> Vec<f64>
where
    F: Fn(&Vec<f64>, &Vec<f64>, usize) -> f64,
    F: Send + Copy + 'static,
{
    let thread_num = 16;
    let mut handles: Vec<thread::JoinHandle<()>> = Vec::new();
    let mut result: Arc<Vec<Mutex<f64>>> = Arc::new(
        vec![0.0; y.len()]
            .into_iter()
            .map(|x| Mutex::new(x))
            .collect(),
    );

    for i in 0..thread_num {
        let x = x.clone();
        let y = y.clone();

        let result = Arc::clone(&mut result);
        let handle = thread::spawn(move || {
            for j in x.len() * i / thread_num..x.len() * (i + 1) / thread_num {
                let mut val = result[j].lock().unwrap();
                *val = 2.0 / PI * f(&x, &y, j);  // <- ボトルネック！！！
            }
        });

        handles.push(handle);
    }

    for handle in handles {
        handle.join().unwrap();
    }

    let result = Arc::try_unwrap(result).unwrap();

    let mut output = vec![0.0; y.len()];
    for (i, val) in result.iter().enumerate() {
        output[i] = *val.lock().unwrap();
    }

    output
}

// ここはPythonで使っていたコードから試行錯誤してすこし変えてみたが根本的な解決
// にはなっていない
fn imag2real_helper(x: &Vec<f64>, y: &Vec<f64>, num: usize) -> f64 {
    let diff = x[1] - x[0];
    let base = x[num];

    let result = x
        .iter()
        .zip(y.iter())
        .map(|(x, y)| {
            if *x == base {
                0.0
            } else {
                x * y / (x.powf(2.0) - base.powf(2.0)) * diff
            }
        })
        .sum();

    result
}

これらのコードで速度を比較してみました。比較した関数は,

• pyfunctionを付けてPython側で使えるようにした imag2real (一番上のコード)
• Rust用に書いた imag2real (一番下のコード)

です。速度を計測したところ, Rustで使うように書いた imag2realが2桁ほど遅くなっていました。細かく時間を計測してみると imag2real_helper で全体の9割くらいの時間を使っていたことがわかったのですが、どこを改善すれば速くなるのかわかりません。

よろしくお願いいたします。

参考情報

Python拡張として書いたRustコード:
https://github.com/Hayashi-Yudai/pykk

Comments

[melian]

昔に書いたという Python スクリプトでは Numpy を使っていましたか？

[Yhaya]

いえ、RustオンリーでPython拡張を書いておりまして、その拡張機能で作ったメソッドをPythonで呼び出したときの速度を、質問に載せましたコードと比較すると２桁ほど遅いという問題が起きています。

Rustで作ったPython拡張は https://github.com/Hayashi-Yudai/pykk に公開していますので、この問題の解決の参考になるかわかりませんが、よろしかったら見てください。

[melian]

ありがとうございます。読ませていただきます。

[dalance]

コピペしたコードとのことで、2つのコードはほとんど同じですので
この2つを比べて原因を究明するのは難しそうです。
同じコードで速度が違うということは、このコードの外の環境による可能性が高いと思います。
実際に実行して速度比較が可能なリポジトリをご用意いただくのがいいかもしれません。

[tatsuya6502]

何点か質問があるのと、確認してほしいことがあります。

コード内の環境に関係するもの
・imag2real_helper() の内容が kk.rs と一番下のコードで変わってますが、なぜでしょうか？
・一番下のコードの imag2real_helper() の内容を kk.rs と同じにすると、性能差はどうなりますか？

コード外の環境に関係するもの
・Python拡張ですが、PyPIで公開されているWindows版（.pyd）を使っているのでしょうか？
・もしそうなら、WSL2のRust 1.55.0でLinux版（.so）をビルドしてPythonから実行すると性能差はどうなりますか？

[Yhaya]

imag2real_helperは上のように元の状態から変えることで少し高速化できています。Linux版のビルドは、これから試してみます

[dalance]

Python版はWindowsネイティブ環境、Rustのみ版はWSL2ということですか？
その場合、実行中のWindowsから見たCPU使用率はどうでしょう？
どちらも16スレッド分使われていますか?

[Yhaya]

調べていたら、Pythonライブラリとしたコードで、Python内で使うところでバグがあったことが分かりました。修正したところ、どちらのコードも同じくらいの速さであることが分かりました。お騒がせしまして申し訳ありませんでした

[tatsuya6502]

ご報告ありがとうございます。解決してよかったです！

Answer 1

質問に書いたコードに間違いはなく、むしろこれを使うコード内にバグあったせいだと分かりました。お騒がせして申し訳ありません

Answer 2

https://qiita.com/tatsuya6502/items/d50e4b131130aa5b5ab6
古い情報ですが、Linux環境においてRustのf64powfの性能はかなり低い傾向があるのかもしれません。

Answer 3

実行手順の情報がないので詳細はわからないのですが、もしかするとデバッグビルドで実行していないでしょうか？

cargo build や cargo run はデフォルトでデバッグビルドを行います。
リリースビルドにするには cargo build --release や cargo run --release を使います。

Comments

[Yhaya]

回答ありがとうございます。リリースビルドであることは確認しております。

Answer 4

Python側の呼び出しコード、簡単なデータセットなどがよく分からないのですが、
cargo buildする時に--releaseオプションなどの最適化オプションを付け忘れていそうな気がします。

Comments

[Yhaya]

それは確認してみたのですが、ちゃんと--releaseはつけてました