Skip to content

Commit

Permalink
Clean code and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@cgouert
cgouert committed Nov 7, 2023
1 parent b952165 commit f52f967
Show file tree
Hide file tree
Showing 3 changed files with 96 additions and 656 deletions.
228 changes: 51 additions & 177 deletions src/bin/helm.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ use tfhe::{
boolean::gen_keys, generate_keys, shortint::parameters::PARAM_MESSAGE_1_CARRY_1_KS_PBS,
ConfigBuilder,
};
use rand::RngCore;
use rand::rngs::OsRng;

fn main() {
ascii::print_art();
Expand Down Expand Up @@ -134,48 +136,39 @@ fn main() {
{
// Gate mode (GPU)
let mut start = Instant::now();
// let (lwe_dim, _, glwe_dim, poly_size) = (
// LweDimension(722),
// LweDimension(722),
// GlweDimension(1),
// PolynomialSize(512),
// );
// let stddev_glwe: f64 = 0.00000004990272175010415;
// let noise = Variance(stddev_glwe.powf(2.0));
// let (dec_lc, dec_bl) = (DecompositionLevelCount(3), DecompositionBaseLog(6));
// let (ks_lc, ks_bl) = (DecompositionLevelCount(4), DecompositionBaseLog(3));

let (lwe_dim, _, glwe_dim, poly_size) = (
LweDimension(630),
let (lwe_dim, glwe_dim, poly_size) = (
LweDimension(630),
GlweDimension(1),
PolynomialSize(1024),
);
let stddev_glwe: f64 = 0.00000002980232238769531;
let stddev_glwe = 0.00000002980232238769531_f64;
let noise = Variance(stddev_glwe.powf(2.0));
let (dec_lc, dec_bl) = (DecompositionLevelCount(3), DecompositionBaseLog(7));
let (ks_lc, ks_bl) = (DecompositionLevelCount(8), DecompositionBaseLog(2));

// Create random seed
let mut random_bytes = [0; 16];
OsRng.fill_bytes(&mut random_bytes);
let random_u128 = u128::from_be_bytes(random_bytes);

// Create the necessary engines
// Here we need to create a secret to give to the unix seeder, but we skip the actual secret creation
const UNSAFE_SECRET: u128 = 0;
let mut default_engine =
DefaultEngine::new(Box::new(UnixSeeder::new(UNSAFE_SECRET))).unwrap();
DefaultEngine::new(Box::new(UnixSeeder::new(random_u128))).unwrap();
let mut parallel_engine =
DefaultParallelEngine::new(Box::new(UnixSeeder::new(UNSAFE_SECRET)))
DefaultParallelEngine::new(Box::new(UnixSeeder::new(random_u128)))
.unwrap();
let mut cuda_engine = CudaEngine::new(()).unwrap();

// Generate the keys
let h_input_key: LweSecretKey32 =
let h_input_key =
default_engine.generate_new_lwe_secret_key(lwe_dim).unwrap();
let h_lut_key: GlweSecretKey32 = default_engine
.generate_new_glwe_secret_key(glwe_dim, poly_size)
.unwrap();
let h_interm_sk: LweSecretKey32 = default_engine
let h_interm_sk = default_engine
.transform_glwe_secret_key_to_lwe_secret_key(h_lut_key.clone())
.unwrap();
let h_keyswitch_key: LweKeyswitchKey32 = default_engine
let h_keyswitch_key = default_engine
.generate_new_lwe_keyswitch_key(
&h_interm_sk,
&h_input_key,
Expand All @@ -185,7 +178,7 @@ fn main() {
)
.unwrap();
// create a BSK with multithreading
let h_bootstrap_key: LweBootstrapKey32 = parallel_engine
let h_bootstrap_key = parallel_engine
.generate_new_lwe_bootstrap_key(
&h_input_key,
&h_lut_key,
Expand All @@ -194,10 +187,10 @@ fn main() {
noise,
)
.unwrap();
let d_fourier_bsk: CudaFourierLweBootstrapKey32 = cuda_engine
let d_fourier_bsk = cuda_engine
.convert_lwe_bootstrap_key(&h_bootstrap_key)
.unwrap();
let d_fourier_ksk: CudaLweKeyswitchKey32 = cuda_engine
let d_fourier_ksk = cuda_engine
.convert_lwe_keyswitch_key(&h_keyswitch_key)
.unwrap();

Expand Down Expand Up @@ -293,166 +286,47 @@ fn main() {
color::Fg(color::LightYellow),
color::Fg(color::Reset)
);
// LUT mode
let mut start = Instant::now();
let (client_key, server_key) =
tfhe::shortint::gen_keys(PARAM_MESSAGE_1_CARRY_1_KS_PBS); // single bit ctxt
let mut circuit = circuit::LutCircuit::new(client_key, server_key, circuit_ptxt);
println!("KeyGen done in {} seconds.", start.elapsed().as_secs_f64());

// Client encrypts their inputs
start = Instant::now();
let mut enc_wire_map =
EvalCircuit::encrypt_inputs(&mut circuit, &wire_set, &input_wire_map);
println!(
"Encryption done in {} seconds.",
start.elapsed().as_secs_f64()
);

if gpu_eval {
#[cfg(feature = "gpu")]
{
// LUT mode
let mut start = Instant::now();
let (lwe_dim, _, glwe_dim, poly_size) = (
LweDimension(656),
LweDimension(512),
GlweDimension(1),
PolynomialSize(512),
);
let stddev_glwe: f64 = 0.00000004053919869756513;
// let stddev_glwe: f64 = 0.0;
let noise = Variance(stddev_glwe.powf(2.0));
let (dec_lc, dec_bl) = (DecompositionLevelCount(2), DecompositionBaseLog(8));
let (ks_lc, ks_bl) = (DecompositionLevelCount(4), DecompositionBaseLog(3));

// Create the necessary engines
// Here we need to create a secret to give to the unix seeder, but we skip the actual secret creation
const UNSAFE_SECRET: u128 = 0;
let mut default_engine =
DefaultEngine::new(Box::new(UnixSeeder::new(UNSAFE_SECRET))).unwrap();
let mut parallel_engine =
DefaultParallelEngine::new(Box::new(UnixSeeder::new(UNSAFE_SECRET)))
.unwrap();
let mut cuda_engine = CudaEngine::new(()).unwrap();
let cuda_amortized_engine = AmortizedCudaEngine::new(()).unwrap();

let h_input_key: LweSecretKey64 =
default_engine.generate_new_lwe_secret_key(lwe_dim).unwrap();

let h_lut_key: GlweSecretKey64 = default_engine
.generate_new_glwe_secret_key(glwe_dim, poly_size)
.unwrap();

let h_interm_sk: LweSecretKey64 = default_engine
.transform_glwe_secret_key_to_lwe_secret_key(h_lut_key.clone())
.unwrap();

let h_keyswitch_key: LweKeyswitchKey64 = default_engine
.generate_new_lwe_keyswitch_key(
&h_interm_sk,
&h_input_key,
ks_lc,
ks_bl,
noise,
)
.unwrap();

// create a BSK with multithreading
let h_bootstrap_key: LweBootstrapKey64 = parallel_engine
.generate_new_lwe_bootstrap_key(
&h_input_key,
&h_lut_key,
dec_bl,
dec_lc,
noise,
)
.unwrap();

let d_fourier_bsk: CudaFourierLweBootstrapKey64 = cuda_engine
.convert_lwe_bootstrap_key(&h_bootstrap_key)
.unwrap();

let d_fourier_ksk: CudaLweKeyswitchKey64 = cuda_engine
.convert_lwe_keyswitch_key(&h_keyswitch_key)
.unwrap();

let mut circuit = circuit::LutCircuitCuda::new(
circuit_ptxt,
default_engine,
cuda_engine,
cuda_amortized_engine,
h_input_key,
h_interm_sk,
d_fourier_bsk,
d_fourier_ksk,
lwe_dim,
glwe_dim,
noise,
poly_size,
);

println!("KeyGen done in {} seconds.", start.elapsed().as_secs_f64());

// Client encrypts their inputs
start = Instant::now();
let mut enc_wire_map =
EvalCircuit::encrypt_inputs(&mut circuit, &wire_set, &input_wire_map);
println!(
"Encryption done in {} seconds.",
start.elapsed().as_secs_f64()
);
start = Instant::now();
enc_wire_map = EvalCircuit::evaluate_encrypted(
&mut circuit,
&enc_wire_map,
1,
arithmetic_type,
);
println!(
"GPU Evaluation done in {} seconds.\n",
start.elapsed().as_secs_f64()
);

// Client decrypts the output of the circuit
start = Instant::now();
println!("Encrypted Evaluation:");
let decrypted_outputs =
EvalCircuit::decrypt_outputs(&mut circuit, &enc_wire_map, verbose);
verilog_parser::write_output_wires(outputs_filename, &decrypted_outputs);
println!(
"Decryption done in {} seconds.",
start.elapsed().as_secs_f64()
);
}
} else {
// LUT mode
let mut start = Instant::now();
let (client_key, server_key) =
tfhe::shortint::gen_keys(PARAM_MESSAGE_1_CARRY_1_KS_PBS); // single bit ctxt
let mut circuit = circuit::LutCircuit::new(client_key, server_key, circuit_ptxt);
println!("KeyGen done in {} seconds.", start.elapsed().as_secs_f64());

// Client encrypts their inputs
for cycle in 0..num_cycles {
start = Instant::now();
let mut enc_wire_map =
EvalCircuit::encrypt_inputs(&mut circuit, &wire_set, &input_wire_map);
println!(
"Encryption done in {} seconds.",
start.elapsed().as_secs_f64()
enc_wire_map = EvalCircuit::evaluate_encrypted(
&mut circuit,
&enc_wire_map,
1,
arithmetic_type,
);

for cycle in 0..num_cycles {
start = Instant::now();
enc_wire_map = EvalCircuit::evaluate_encrypted(
&mut circuit,
&enc_wire_map,
1,
arithmetic_type,
);
println!(
"Cycle {}) Evaluation done in {} seconds.\n",
cycle,
start.elapsed().as_secs_f64()
);
}

// Client decrypts the output of the circuit
start = Instant::now();
println!("Encrypted Evaluation:");
let decrypted_outputs =
EvalCircuit::decrypt_outputs(&mut circuit, &enc_wire_map, verbose);
verilog_parser::write_output_wires(outputs_filename, &decrypted_outputs);
println!(
"Decryption done in {} seconds.",
"Cycle {}) Evaluation done in {} seconds.\n",
cycle,
start.elapsed().as_secs_f64()
);
}

// Client decrypts the output of the circuit
start = Instant::now();
println!("Encrypted Evaluation:");
let decrypted_outputs =
EvalCircuit::decrypt_outputs(&mut circuit, &enc_wire_map, verbose);
verilog_parser::write_output_wires(outputs_filename, &decrypted_outputs);
println!(
"Decryption done in {} seconds.",
start.elapsed().as_secs_f64()
);
}
}
println!();
Expand Down
Loading

0 comments on commit f52f967

Please sign in to comment.