J’ai optimisé un bout de code avec Rust et j’ai gagné un facteur ~670 dans les performances.

• la version initiale en python prend 30s pour 1 million d’itérations
• la version optimisée en Rust prend 0.4s pour 10 million d’itérations

Summary
  './target/release/kero_fold_par_iter 1000000' ran
    7.01 ± 0.28 times faster than './target/release/kero_simple_port 1000000'
  667.27 ± 25.77 times faster than 'python main.py'

Nous allons voir dans cet article comment j’ai optimisé le code, en pilotant les améliorations grâce à des outils.

J’ai fait plusieurs choses:

• écriture d’une première version Python
• réécrire le code en Rust
• benchmarker avec Criterion
• paralléliser avec Rayon
• micro-optimiser avec perf+flamegraph
• visualisation des résultats avec plotters

Le problème et première version en Python

Mon problème : calculer des probabilités pour un bout du jeu de plateau Kero. Pour cela, j’ai écrit une simulation de Monte Carlo.

Le principe est simple:

• on simule plein de fois un processus aléatoire
• la loi des grands nombres nous dit en gros que si on fait suffisament de simulations, la distribution des résultats devrait être représentative du résultat analytique.

Concrètement, si on prend l’exemple d’un dé à 6 faces: Si vous lancez beaucoup de fois un dé à 6 faces (disons, un million), vous devriez avoir à peu prêt autant de fois des 1, des 2… que des 6. De cette distribution des résultats, on peut déduire des probabilités, par exemple celle d’obtenir un dé de valeur 6.

Dans cet exemple très simple on pourrait calculer le résultat analytique facilement, mais il y a des cas où ca n’est pas possible, ou bien où il est plus simple de faire des estimations. Exemple pour un problème précédent: https://www.keiruaprod.fr/blog/2020/12/10/better-odds-against-my-3yo-kid.html

J’ai donc écrit une première version:

def simulate_refill(nb_dices=8):
	# on simule des lancers de dés, le code est peu important
	# …
	return nb_rolls


dices_rolls = defaultdict(int)
for i in range(1_000_000):
	nb_rolls = simulate_refill()
	dices_rolls[nb_rolls] += 1

Cette première version prend ~30s:

$ time python main.py
real	0m29,254s
user	0m29,219s
sys	0m0,031s

À noter que time est une très mauvaise manière de calculer les performances, mais ça suffit à avoir un ordre d’idée au départ.

À noter aussi qu’à ce stade j’avais la réponse à ma question sur Kero, mais j’étais curieux de voir jusqu’où Rust pouvait raisonnablement nous amener.

Réécriture en Rust

J’ai réécrit la boucle principale en Rust, en stockant les résultats dans une HashMap de la librairie standard. Le code est très proche, il faut inclure la crate rand et implémenter simulate_refill mais ce n’est pas le sujet de l’article.

pub fn generate_rolls_counts(nb_simulations: u32) -> HashMap<u32, u32> {
    let mut rng = rand::thread_rng();
    let mut dice_rolls: HashMap<u32, u32> = HashMap::new();
    for _ in 0..nb_simulations {
        let rolls = simulate_refill(&mut rng, 8);
        *dice_rolls.entry(rolls).or_insert(0) += 1;
    }
    dice_rolls
}

Grossièrement avec time, on est déjà plus rapide d’un facteur 100.

On peut aussi écrire cette fonction différement, avec une approche fonctionnelle en utilisant fold de la librairie standard.

pub fn generate_rolls_counts_fold(nb_simulations: u32) -> HashMap<u32, u32> {
    let mut rng = rand::thread_rng();
    let dice_rolls: HashMap<u32, u32> = (0..nb_simulations)
        .into_iter()
        .fold(HashMap::new(), |mut acc, _|{
            let rolls = simulate_refill(&mut rng, 8);
            *acc.entry(rolls).or_insert(0) += 1;
            acc
        });
    dice_rolls
}

Le résulat est le même, le style est différent (itératif vs fonctionnel), quelle est la fonction la plus rapide ?

Pour y répondre correctement, il faut benchmarker.

Benchmarking

Il y a plusieurs problèmes au fait de faire des benchmarks naifs avec time: la gestion du cache (dans quel état est-il? chaud, pas chaud?), la precision des résulats (la seconde), le nombre de résultats (1 seul).

Une solution plus pérenne c’est d’utiliser un outil dédié. Rust dispose d’une crate, criterion, qui sert à ça. La doc est très bonne. Une des fonctionnalités: comparer deux implémentations.

Je passe l’installation et la configuration, venons en directement au fichier de benchmark:

// benches/simulation.rs
use criterion::{criterion_group, criterion_main, Criterion, BenchmarkId};
use kero::*;

fn bench_generate_rolls(c: &mut Criterion) {
    let mut group = c.benchmark_group("Generate rolls");
    for i in [100_000].iter() {
        group.bench_with_input(BenchmarkId::new("Linear", i), i, 
            |b, i| b.iter(|| generate_rolls_counts(*i)));
        group.bench_with_input(BenchmarkId::new("Fold", i), i, 
            |b, i| b.iter(|| generate_rolls_counts_fold(*i)));
    }
    group.finish();
}

criterion_group!(benches, bench_generate_rolls);
criterion_main!(benches);

Avec ca, un cargo bench nous permet

Parallèliser

Le problème de nos deux versions, c’est que tout se fait dans un même thread, séquentiellement.

$ cargo add rayon
    Updating 'https://github.com/rust-lang/crates.io-index' index
      Adding rayon v1.5.3 to dependencies

On ne peut pas simplement remplacer iter par par_iter comme c’est parfois le cas : chaque thread se retrouverait à écrire dans la même HashMap, ce qui pose des problèmes de concurrence d’accès.

Une solution: chaque thread va créer une Hashmap, puis elles vont être combinées avec reduce pour retourner une unique Hashmap.

pub fn generate_rolls_counts_fold_par_iter(nb_simulations: u32) -> HashMap<u32, u32> { 
   let dice_rolls: HashMap<u32, u32> = (0..nb_simulations)
        .into_par_iter()
        .fold(HashMap::new,|mut acc, _|{
            let mut rng = rand::thread_rng();
            let rolls = simulate_refill(&mut rng, 8);
            *acc.entry(rolls).or_insert(0) += 1;
            acc
        })
        .reduce_with(|mut m1, m2| {
            for (k, v) in m2 {
                *m1.entry(k).or_default() += v;
            }
            m1
        })
        .unwrap();
    dice_rolls
}

flamegraph

Si on veut aller plus loin dans le détail, on peut utiliser un flamegraph.

cargo install flamegraph
cargo flamegraph --bin kero_fold_par_iter

Personnellement à ce stade j’étais déjà content.

Benchmarking des binaires avec hyperfine

hyperfine propose une autre approche du benchmarking.

Dans des scénarios plus compliqués, il peut être pertinent de faire les benchmarks sur les binaires complets, et pas juste sur les fonctions de la librairie. Cela peut permettre d’avoir des tests plus réalistes, en comparant les temps d’exécution de deux versions. Ca permet plusieurs choses, comme comparer des programmes dans différents langages. Il est aussi possible que l’on ait pas accès au code source de l’autre version.

Ici, cela nous permet de comparer la version initiale en python avec différentes améliorations. On peut voir que le gain est d’environ un facteur 667 entre la version initiale et une version parallélisée en Rust.

$ cargo build --release
$ hyperfine "python main.py 1000000" "./target/release/kero_simple_port 1000000" "./target/release/kero_fold_par_iter 1000000"
Benchmark 1: python main.py
 ⠙ Current estimate: 33.236 s ████████████████████████████████████████████████████████████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░  Time (mean ± σ):     33.010 s ±  0.764 s    [User: 32.990 s, System: 0.018 s]
  Range (min … max):   31.800 s … 34.150 s    10 runs
 
Benchmark 2: ./target/release/kero_simple_port 1000000
  Time (mean ± σ):     346.8 ms ±   8.8 ms    [User: 340.7 ms, System: 6.1 ms]
  Range (min … max):   338.3 ms … 369.8 ms    10 runs
 
Benchmark 3: ./target/release/kero_fold_par_iter 1000000
  Time (mean ± σ):      49.5 ms ±   1.5 ms    [User: 574.4 ms, System: 14.1 ms]
  Range (min … max):    45.2 ms …  53.1 ms    58 runs
 
Summary
  './target/release/kero_fold_par_iter 1000000' ran
    7.01 ± 0.28 times faster than './target/release/kero_simple_port 1000000'
  667.27 ± 25.77 times faster than 'python main.py'