Add dpa helper function and benchmark

Cargo.toml

@@ -36,3 +36,7 @@ harness = false
 [[bench]]
 name = "snr"
 harness = false
+
+[[bench]]
+name = "dpa"
+harness = false

benches/dpa.rs

@@ -0,0 +1,70 @@
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use muscat::dpa::{dpa, Dpa};
+use muscat::leakage::sbox;
+use ndarray::{Array1, Array2};
+use ndarray_rand::rand::{rngs::StdRng, SeedableRng};
+use ndarray_rand::rand_distr::Uniform;
+use ndarray_rand::RandomExt;
+
+fn selection_function(metadata: Array1<u8>, guess: usize) -> usize {
+    sbox(metadata[1] ^ guess as u8).into()
+}
+
+fn dpa_sequential(leakages: &Array2<f32>, plaintexts: &Array2<u8>) -> Dpa<Array1<u8>> {
+    let mut dpa = Dpa::new(leakages.shape()[1], 256, selection_function);
+
+    for i in 0..leakages.shape()[0] {
+        dpa.update(leakages.row(i), plaintexts.row(i).to_owned());
+    }
+
+    dpa.finalize();
+
+    dpa
+}
+
+fn dpa_parallel(leakages: &Array2<f32>, plaintexts: &Array2<u8>) -> Dpa<Array1<u8>> {
+    dpa(
+        leakages.view(),
+        plaintexts
+            .rows()
+            .into_iter()
+            .map(|x| x.to_owned())
+            .collect::<Array1<Array1<u8>>>()
+            .view(),
+        256,
+        selection_function,
+        500,
+    )
+}
+
+fn bench_dpa(c: &mut Criterion) {
+    // Seed rng to get the same output each run
+    let mut rng = StdRng::seed_from_u64(0);
+
+    let mut group = c.benchmark_group("dpa");
+
+    group.measurement_time(std::time::Duration::from_secs(60));
+
+    for nb_traces in [1000, 2000, 5000].into_iter() {
+        let leakages = Array2::random_using((nb_traces, 5000), Uniform::new(-2., 2.), &mut rng);
+        let plaintexts =
+            Array2::random_using((nb_traces, 16), Uniform::new_inclusive(0, 255), &mut rng);
+
+        group.bench_with_input(
+            BenchmarkId::new("sequential", nb_traces),
+            &(&leakages, &plaintexts),
+            |b, (leakages, plaintexts)| b.iter(|| dpa_sequential(leakages, plaintexts)),
+        );
+
+        group.bench_with_input(
+            BenchmarkId::new("parallel", nb_traces),
+            &(&leakages, &plaintexts),
+            |b, (leakages, plaintexts)| b.iter(|| dpa_parallel(leakages, plaintexts)),
+        );
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_dpa);
+criterion_main!(benches);

examples/dpa.rs

@@ -1,9 +1,9 @@
 use anyhow::Result;
 use indicatif::ProgressIterator;
-use muscat::dpa::*;
+use muscat::dpa::Dpa;
 use muscat::leakage::sbox;
 use muscat::util::read_array2_from_npy_file;
-use ndarray::*;
+use ndarray::{s, Array1, Array2};
 use rayon::iter::{ParallelBridge, ParallelIterator};
 
 // traces format

src/cpa.rs

@@ -41,7 +41,7 @@ where
     })
     .reduce(
         || Cpa::new(leakages.shape()[1], guess_range, target_byte, leakage_func),
-        |a: Cpa, b| a + b,
+        |a, b| a + b,
     );
 
     cpa.finalize();

src/dpa.rs

@@ -1,9 +1,45 @@
 use ndarray::{concatenate, Array1, Array2, ArrayView1, ArrayView2, Axis};
-use std::ops::Add;
+use rayon::iter::{ParallelBridge, ParallelIterator};
+use std::{iter::zip, ops::Add};
 
 use crate::util::max_per_row;
 
-pub struct Dpa<T> {
+pub fn dpa<M, T>(
+    leakages: ArrayView2<T>,
+    metadata: ArrayView1<M>,
+    guess_range: usize,
+    leakage_func: fn(M, usize) -> usize,
+    chunk_size: usize,
+) -> Dpa<M>
+where
+    T: Into<f32> + Copy + Sync,
+    M: Clone + Sync,
+{
+    let mut dpa = zip(
+        leakages.axis_chunks_iter(Axis(0), chunk_size),
+        metadata.axis_chunks_iter(Axis(0), chunk_size),
+    )
+    .par_bridge()
+    .map(|(leakages_chunk, metadata_chunk)| {
+        let mut dpa = Dpa::new(leakages.shape()[1], guess_range, leakage_func);
+
+        for i in 0..leakages_chunk.shape()[0] {
+            dpa.update(leakages_chunk.row(i), metadata_chunk[i].clone());
+        }
+
+        dpa
+    })
+    .reduce(
+        || Dpa::new(leakages.shape()[1], guess_range, leakage_func),
+        |a, b| a + b,
+    );
+
+    dpa.finalize();
+
+    dpa
+}
+
+pub struct Dpa<M> {
     /// Number of samples per trace
     len_samples: usize,
     /// Guess range upper excluded bound
@@ -21,7 +57,7 @@ pub struct Dpa<T> {
     rank_slice: Array2<f32>,
     rank_traces: usize, // Number of traces to calculate succes rate
     /// Selection function
-    leakage_func: fn(T, usize) -> usize,
+    leakage_func: fn(M, usize) -> usize,
     /// Number of traces processed
     len_leakages: usize,
 }
@@ -30,8 +66,8 @@ pub struct Dpa<T> {
 https://paulkocher.com/doc/DifferentialPowerAnalysis.pdf
 https://web.mit.edu/6.857/OldStuff/Fall03/ref/kocher-DPATechInfo.pdf */
 
-impl<T: Clone> Dpa<T> {
-    pub fn new(size: usize, guess_range: usize, f: fn(T, usize) -> usize) -> Self {
+impl<M: Clone> Dpa<M> {
+    pub fn new(size: usize, guess_range: usize, f: fn(M, usize) -> usize) -> Self {
         Self {
             len_samples: size,
             guess_range,
@@ -50,9 +86,9 @@ impl<T: Clone> Dpa<T> {
 
     /// # Panics
     /// Panic in debug if `trace.shape()[0] != self.len_samples`.
-    pub fn update<U>(&mut self, trace: ArrayView1<U>, metadata: T)
+    pub fn update<T>(&mut self, trace: ArrayView1<T>, metadata: M)
     where
-        U: Into<f32> + Copy,
+        T: Into<f32> + Copy,
     {
         debug_assert_eq!(trace.shape()[0], self.len_samples);
 
@@ -81,9 +117,9 @@ impl<T: Clone> Dpa<T> {
         self.len_leakages += 1;
     }
 
-    pub fn update_success<U>(&mut self, trace_batch: ArrayView1<U>, plaintext_batch: T)
+    pub fn update_success<T>(&mut self, trace_batch: ArrayView1<T>, plaintext_batch: M)
     where
-        U: Into<f32> + Copy,
+        T: Into<f32> + Copy,
     {
         /* This function updates the main arrays of the DPA for the success rate*/
         self.update(trace_batch, plaintext_batch);
@@ -160,7 +196,7 @@ impl<T: Clone> Dpa<T> {
     }
 }
 
-impl<T> Add for Dpa<T> {
+impl<M> Add for Dpa<M> {
     type Output = Self;
 
     fn add(self, rhs: Self) -> Self::Output {