Fix slow preproc loading

alexandrainst · Aug 22, 2024 · fc7182b · fc7182b
1 parent 9860e8b
commit fc7182b
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 24 deletions.
diff --git a/src/schemes/spdz/mod.rs b/src/schemes/spdz/mod.rs
@@ -1641,7 +1641,7 @@ mod test {
         let mut files = [tempfile::tempfile().unwrap(), tempfile::tempfile().unwrap()];
         let known_to_each = vec![1, 2];
         let number_of_triplets = 2;
-        preprocessing::write_preproc_to_file(
+        preprocessing::write_context(
             &mut files,
             known_to_each,
             number_of_triplets,
@@ -1650,8 +1650,8 @@ mod test {
         .unwrap();
         files[0].rewind().unwrap();
         files[1].rewind().unwrap();
-        let p1_context: SpdzContext<F> = preprocessing::read_preproc_from_file(&mut files[0]);
-        let p2_context: SpdzContext<F> = preprocessing::read_preproc_from_file(&mut files[1]);
+        let p1_context: SpdzContext<F> = preprocessing::load_context(&mut files[0]);
+        let p2_context: SpdzContext<F> = preprocessing::load_context(&mut files[1]);
         // unpacking
         let p1_params = p1_context.params;
         let p2_params = p2_context.params;

diff --git a/src/schemes/spdz/preprocessing.rs b/src/schemes/spdz/preprocessing.rs
@@ -7,11 +7,13 @@ use crate::{
 use bincode;
 use ff::PrimeField;
 use rand::SeedableRng;
+use serde::{de::DeserializeOwned, Serialize};
 use std::{
     error::Error,
     fmt,
     fs::File,
     io::{self, Write},
+    path::Path,
 };
 
 #[derive(Debug, Clone, PartialEq)]
@@ -93,7 +95,7 @@ pub struct SecretValues<F> {
     pub mac_key: F,
 }
 
-pub fn write_preproc_to_file<F: PrimeField + serde::Serialize + serde::de::DeserializeOwned>(
+pub fn write_context<F: PrimeField + Serialize + DeserializeOwned>(
     files: &mut [File],
     known_to_each: Vec<usize>,
     number_of_triplets: usize,
@@ -113,14 +115,23 @@ pub fn write_preproc_to_file<F: PrimeField + serde::Serialize + serde::de::Deser
     Ok(())
 }
 
-pub fn read_preproc_from_file<F: PrimeField + serde::Serialize + serde::de::DeserializeOwned>(
+pub fn load_context<F: PrimeField + Serialize + DeserializeOwned>(
     file: &mut File,
 ) -> SpdzContext<F> {
     // TODO: return Result instead.
-    bincode::deserialize_from(file).unwrap()
+    let buffered = std::io::BufReader::new(file);
+    bincode::deserialize_from(buffered).unwrap()
 }
 
-pub fn dealer_preproc<F: PrimeField + serde::Serialize + serde::de::DeserializeOwned>(
+pub async fn load_context_async<F: PrimeField + Serialize + DeserializeOwned>(
+    file: &Path,
+) -> SpdzContext<F> {
+    let contents = tokio::fs::read(file).await.unwrap();
+    // TODO: return Result instead.
+    bincode::deserialize_from(&*contents).unwrap()
+}
+
+pub fn dealer_preproc<F: PrimeField + Serialize + DeserializeOwned>(
     mut rng: impl rand::Rng,
     known_to_each: Vec<usize>,
     number_of_triplets: usize,
@@ -247,13 +258,13 @@ pub fn dealer_preproc<F: PrimeField + serde::Serialize + serde::de::DeserializeO
     (contexts, secret_values)
 }
 
-impl<F: PrimeField + serde::Serialize + serde::de::DeserializeOwned> SpdzContext<F> {
+impl<F: PrimeField + Serialize + DeserializeOwned> SpdzContext<F> {
     fn empty(number_of_parties: usize, mac_key_share: F, who_am_i: Id) -> Self {
         generate_empty_context(number_of_parties, mac_key_share, who_am_i)
     }
 
     pub fn from_file(mut file: File) -> Result<Self, io::Error> {
-        Ok(read_preproc_from_file(&mut file))
+        Ok(load_context(&mut file))
     }
 }
 

diff --git a/wecare/Cargo.toml b/wecare/Cargo.toml
@@ -21,3 +21,7 @@ tempfile = "3.10.1"
 [[bench]]
 name = "spdz-25519"
 harness = false
+
+[[bench]]
+name = "shamir-25519"
+harness = false
diff --git a/wecare/benches/spdz-25519.rs b/wecare/benches/spdz-25519.rs
@@ -1,24 +1,29 @@
 use criterion::{criterion_group, criterion_main, Criterion};
-use std::time::Duration;
+use std::time;
 use std::{fs::File, hint::black_box, io::Seek, thread};
+use std::{io::Write, time::Duration};
 use wecare::*;
 
 fn precompute(n: usize) -> (File, File) {
+    let clock = time::Instant::now();
     print!("\nPrecomputing...");
+    let _ = std::io::stdout().flush();
     let ctx1 = tempfile::tempfile().unwrap();
     let ctx2 = tempfile::tempfile().unwrap();
     let mut files = [ctx1, ctx2];
-    do_preproc(&mut files, vec![n, n], false);
+    do_preproc(&mut files, &[n, n], false);
     let [mut ctx1, mut ctx2] = files;
     ctx1.rewind().unwrap();
     ctx2.rewind().unwrap();
-    println!(" Complete!");
+    println!(" Complete! (took {:#?})", clock.elapsed());
     (ctx1, ctx2)
 }
 
 fn build_spdz_engines() -> (Engine, Engine) {
-    let (mut ctx1, mut ctx2) = precompute(2000000);
-    print!("\nSetting up engines...");
+    let (mut ctx1, mut ctx2) = precompute(10000000);
+    let clock = time::Instant::now();
+    print!("Setting up engines...");
+    let _ = std::io::stdout().flush();
     let (e1, e2) = thread::scope(|scope| {
         let e2 = scope.spawn(|| {
             Engine::setup("127.0.0.1:1234")
@@ -37,7 +42,7 @@ fn build_spdz_engines() -> (Engine, Engine) {
         });
         (e1.join().unwrap(), e2.join().unwrap())
     });
-    println!(" Complete!");
+    println!(" Complete! (took {:#?})", clock.elapsed());
     (e1, e2)
 }
 

diff --git a/wecare/src/lib.rs b/wecare/src/lib.rs
@@ -136,18 +136,16 @@ impl std::fmt::Display for MpcError {
 
 impl std::error::Error for MpcError {}
 
-pub fn do_preproc(files: &mut [File], number_of_shares: Vec<usize>, use_32: bool) {
+pub fn do_preproc(files: &mut [File], number_of_shares: &[usize], use_32: bool) {
     assert_eq!(files.len(), number_of_shares.len());
     let known_to_each = vec![number_of_shares[0], number_of_shares[1]];
     let number_of_triplets = 0;
     if use_32 {
         let num = S32::from_f64(0.0);
-        preprocessing::write_preproc_to_file(files, known_to_each, number_of_triplets, num)
-            .unwrap();
+        preprocessing::write_context(files, known_to_each, number_of_triplets, num).unwrap();
     } else {
         let num = S25519::from_f64(0.0);
-        preprocessing::write_preproc_to_file(files, known_to_each, number_of_triplets, num)
-            .unwrap();
+        preprocessing::write_context(files, known_to_each, number_of_triplets, num).unwrap();
     }
 }
 
@@ -179,12 +177,12 @@ mod generic {
                 .preprocessed
                 .ok_or(MpcError("No proccesing file found"))?;
             if self.use_32bit_field {
-                let mut context = preprocessing::read_preproc_from_file(file);
+                let mut context = preprocessing::load_context(file);
                 context.params.who_am_i = network.id();
                 let engine = SpdzEngine32::new(network, runtime, context);
                 Ok(AdderEngine::Spdz32(engine))
             } else {
-                let mut context = preprocessing::read_preproc_from_file(file);
+                let mut context = preprocessing::load_context(file);
                 context.params.who_am_i = network.id();
                 let engine = SpdzEngine::new(network, runtime, context);
                 Ok(AdderEngine::Spdz(engine))
@@ -362,7 +360,7 @@ mod test {
         let ctx1 = tempfile::tempfile().unwrap();
         let ctx2 = tempfile::tempfile().unwrap();
         let mut files = [ctx1, ctx2];
-        do_preproc(&mut files, vec![1, 1], false);
+        do_preproc(&mut files, &[1, 1], false);
         let [mut ctx1, mut ctx2] = files;
         ctx1.rewind().unwrap();
         ctx2.rewind().unwrap();
@@ -407,7 +405,7 @@ mod test {
         let ctx1 = tempfile::tempfile().unwrap();
         let ctx2 = tempfile::tempfile().unwrap();
         let mut files = [ctx1, ctx2];
-        do_preproc(&mut files, vec![2, 2], false);
+        do_preproc(&mut files, &[2, 2], false);
         let [mut ctx1, mut ctx2] = files;
         ctx1.rewind().unwrap();
         ctx2.rewind().unwrap();