From bb7216a83232d76bfbb27da2faca9531d440e1ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ga=C3=ABtan=20Cassiers?= Date: Fri, 26 Nov 2021 12:09:41 +0100 Subject: [PATCH 1/2] option to hide progress bar in belief prop --- setup.py | 1 - src/scalib/attacks/sascagraph.py | 5 +++- .../scalib-py/src/belief_propagation.rs | 14 +++++++-- src/scalib_ext/scalib-py/src/lib.rs | 3 +- .../scalib/src/belief_propagation.rs | 29 +++++++++++++------ 5 files changed, 38 insertions(+), 14 deletions(-) diff --git a/setup.py b/setup.py index c1f4b566..486680ae 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,6 @@ # Ensure these are present (in case we are not using PEP-518 compatible build # system). import setuptools_scm -import toml scalib_features = ["pyo3/abi3"] diff --git a/src/scalib/attacks/sascagraph.py b/src/scalib/attacks/sascagraph.py index ad264048..0759aea7 100644 --- a/src/scalib/attacks/sascagraph.py +++ b/src/scalib/attacks/sascagraph.py @@ -223,13 +223,15 @@ def set_table(self, table, values): ) self.tables_[table] = values - def run_bp(self, it): + def run_bp(self, it, progress=False): r"""Runs belief propagation algorithm on the current state of the graph. Parameters ---------- it : int Number of iterations of belief propagation. + progress: bool + Show a progress bar (default: False). """ if self.solved_: raise Exception("Cannot run bp twice on a graph.") @@ -241,6 +243,7 @@ def run_bp(self, it): self.edge_, self.nc_, self.n_, + progress, ) self.solved_ = True diff --git a/src/scalib_ext/scalib-py/src/belief_propagation.rs b/src/scalib_ext/scalib-py/src/belief_propagation.rs index c9890dd0..b20580c8 100644 --- a/src/scalib_ext/scalib-py/src/belief_propagation.rs +++ b/src/scalib_ext/scalib-py/src/belief_propagation.rs @@ -89,6 +89,8 @@ pub fn run_bp( nc: usize, // number of copies in the graph (n_runs) n: usize, + // show a progress bar + progress: bool, ) -> PyResult<()> { // map all python functions to rust ones + generate the mapping in vec_functs_id let functions_rust: Vec = functions @@ -104,8 +106,16 @@ pub fn run_bp( .map(|x| to_var(x.downcast::().unwrap())) .collect(); - scalib::belief_propagation::run_bp(&functions_rust, &mut variables_rust, it, edge, nc, n) - .unwrap(); + scalib::belief_propagation::run_bp( + &functions_rust, + &mut variables_rust, + it, + edge, + nc, + n, + progress, + ) + .unwrap(); variables_rust .iter() diff --git a/src/scalib_ext/scalib-py/src/lib.rs b/src/scalib_ext/scalib-py/src/lib.rs index aa52d4ca..8f8cfa85 100644 --- a/src/scalib_ext/scalib-py/src/lib.rs +++ b/src/scalib_ext/scalib-py/src/lib.rs @@ -43,8 +43,9 @@ fn _scalib_ext(_py: Python, m: &PyModule) -> PyResult<()> { vertex: usize, nc: usize, n: usize, + progress: bool, ) -> PyResult<()> { - belief_propagation::run_bp(py, functions, variables, it, vertex, nc, n) + belief_propagation::run_bp(py, functions, variables, it, vertex, nc, n, progress) } #[pyfn(m, "partial_cp")] diff --git a/src/scalib_ext/scalib/src/belief_propagation.rs b/src/scalib_ext/scalib/src/belief_propagation.rs index 3a4558c2..482451b1 100644 --- a/src/scalib_ext/scalib/src/belief_propagation.rs +++ b/src/scalib_ext/scalib/src/belief_propagation.rs @@ -351,6 +351,8 @@ pub fn run_bp( nc: usize, // number of copies in the graph (n_runs) n: usize, + // show a progress bar + progress: bool, ) -> Result<(), ()> { // Scratch array containing all the edge's messages. let mut edges: Vec> = vec![Array2::::ones((n, nc)); edge]; @@ -386,15 +388,7 @@ pub fn run_bp( } } - // loading bar - let pb = ProgressBar::new(it as u64); - pb.set_style(ProgressStyle::default_spinner().template( - "{msg} {spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] ({pos}/{len}, ETA {eta})", - ) - .on_finish(ProgressFinish::AndClear)); - pb.set_message("Calculating BP..."); - - for _ in (0..it).progress_with(pb) { + let mut bp_iter = || { // This is a technique for runtime borrow-checking: we take reference on all the edges // at once, put them into options, then extract the references out of the options, one // at a time and out-of-order. @@ -422,6 +416,23 @@ pub fn run_bp( }) .collect(); update_variables(&mut edge_for_var, variables); + }; + + if progress { + // loading bar + let pb = ProgressBar::new(it as u64); + pb.set_style(ProgressStyle::default_spinner().template( + "{msg} {spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] ({pos}/{len}, ETA {eta})", + ) + .on_finish(ProgressFinish::AndClear)); + pb.set_message("Calculating BP..."); + for _ in (0..it).progress_with(pb) { + bp_iter(); + } + } else { + for _ in 0..it { + bp_iter(); + } } Ok(()) From 0228c2599dee91bfe2c5ffa09615b5372bfc4d27 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ga=C3=ABtan=20Cassiers?= Date: Fri, 26 Nov 2021 16:24:57 +0100 Subject: [PATCH 2/2] Make rank estimation and probability estimation for LDA run in parallel --- .../scalib-py/src/belief_propagation.rs | 22 ++++++------ src/scalib_ext/scalib-py/src/lda.rs | 2 +- src/scalib_ext/scalib-py/src/lib.rs | 34 +++++++++++-------- 3 files changed, 33 insertions(+), 25 deletions(-) diff --git a/src/scalib_ext/scalib-py/src/belief_propagation.rs b/src/scalib_ext/scalib-py/src/belief_propagation.rs index b20580c8..a8dd65df 100644 --- a/src/scalib_ext/scalib-py/src/belief_propagation.rs +++ b/src/scalib_ext/scalib-py/src/belief_propagation.rs @@ -106,16 +106,18 @@ pub fn run_bp( .map(|x| to_var(x.downcast::().unwrap())) .collect(); - scalib::belief_propagation::run_bp( - &functions_rust, - &mut variables_rust, - it, - edge, - nc, - n, - progress, - ) - .unwrap(); + py.allow_threads(|| { + scalib::belief_propagation::run_bp( + &functions_rust, + &mut variables_rust, + it, + edge, + nc, + n, + progress, + ) + .unwrap(); + }); variables_rust .iter() diff --git a/src/scalib_ext/scalib-py/src/lda.rs b/src/scalib_ext/scalib-py/src/lda.rs index 4128cec7..9501d200 100644 --- a/src/scalib_ext/scalib-py/src/lda.rs +++ b/src/scalib_ext/scalib-py/src/lda.rs @@ -96,7 +96,7 @@ impl LDA { /// x : traces with shape (n,ns) /// return prs with shape (n,nc). Every row corresponds to one probability distribution fn predict_proba<'py>( - &mut self, + &self, py: Python<'py>, x: PyReadonlyArray2, ) -> PyResult<&'py PyArray2> { diff --git a/src/scalib_ext/scalib-py/src/lib.rs b/src/scalib_ext/scalib-py/src/lib.rs index 8f8cfa85..9fc5c108 100644 --- a/src/scalib_ext/scalib-py/src/lib.rs +++ b/src/scalib_ext/scalib-py/src/lib.rs @@ -70,6 +70,7 @@ fn _scalib_ext(_py: Python, m: &PyModule) -> PyResult<()> { #[pyfn(m, "rank_accuracy")] fn rank_accuracy( + py: Python, costs: Vec>, key: Vec, acc: f64, @@ -77,32 +78,37 @@ fn _scalib_ext(_py: Python, m: &PyModule) -> PyResult<()> { method: String, max_nb_bin: usize, ) -> PyResult<(f64, f64, f64)> { - let res = str2method(&method).unwrap_or_else(|s| panic!("{}", s)); - let res = res.rank_accuracy(&costs, &key, acc, merge, max_nb_bin); - match res { - Ok(res) => Ok((res.min, res.est, res.max)), - Err(s) => { - panic!("{}", s); + py.allow_threads(|| { + let res = str2method(&method).unwrap_or_else(|s| panic!("{}", s)); + let res = res.rank_accuracy(&costs, &key, acc, merge, max_nb_bin); + match res { + Ok(res) => Ok((res.min, res.est, res.max)), + Err(s) => { + panic!("{}", s); + } } - } + }) } #[pyfn(m, "rank_nbin")] fn rank_nbin( + py: Python, costs: Vec>, key: Vec, nb_bin: usize, merge: Option, method: String, ) -> PyResult<(f64, f64, f64)> { - let res = str2method(&method).unwrap_or_else(|s| panic!("{}", s)); - let res = res.rank_nbin(&costs, &key, nb_bin, merge); - match res { - Ok(res) => Ok((res.min, res.est, res.max)), - Err(s) => { - panic!("{}", s); + py.allow_threads(|| { + let res = str2method(&method).unwrap_or_else(|s| panic!("{}", s)); + let res = res.rank_nbin(&costs, &key, nb_bin, merge); + match res { + Ok(res) => Ok((res.min, res.est, res.max)), + Err(s) => { + panic!("{}", s); + } } - } + }) } Ok(())