Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

build: move tfhe ops into separate crate #8

Merged
merged 1 commit into from
Aug 21, 2024
Merged

build: move tfhe ops into separate crate #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions fhevm-engine/Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions fhevm-engine/coprocessor/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,11 @@ regex = "1.10.5"
lazy_static = "1.5.0"
clap.workspace = true
lru = "0.12.3"
bincode = "1.3.3"
hex = "0.4"
strum = { version = "0.26", features = ["derive"] }
bigdecimal = "0.4"
fhevm-engine-common = { path = "../fhevm-engine-common" }
strum = { version = "0.26", features = ["derive"] }
bincode = "1.3.3"

[dev-dependencies]
testcontainers = "0.21"
Expand Down
22 changes: 7 additions & 15 deletions fhevm-engine/coprocessor/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,28 +1,24 @@
use fhevm_engine_common::generate_fhe_keys;
use tokio::task::JoinSet;

mod cli;
mod db_queries;
mod server;
#[cfg(test)]
mod tests;
mod tfhe_ops;
mod tfhe_worker;
mod types;
mod utils;

fn main() {

// TODO: remove, just to make sure it works
let _ = fhevm_engine_common::add(5, 5);

let args = crate::cli::parse_args();
assert!(
args.work_items_batch_size < args.tenant_key_cache_size,
"Work items batch size must be less than tenant key cache size"
);

if args.generate_fhe_keys {
generate_fhe_keys();
generate_dump_fhe_keys();
} else {
start_runtime(args, None);
}
Expand Down Expand Up @@ -85,20 +81,16 @@ async fn async_main(
Ok(())
}

fn generate_fhe_keys() {
fn generate_dump_fhe_keys() {
let output_dir = "fhevm-keys";
println!("Generating keys...");
let (client_key, server_key) = tfhe::generate_keys(tfhe::ConfigBuilder::default().build());
let compact_key = tfhe::CompactPublicKey::new(&client_key);
let client_key = bincode::serialize(&client_key).unwrap();
let server_key = bincode::serialize(&server_key).unwrap();
let compact_key = bincode::serialize(&compact_key).unwrap();
let keys = generate_fhe_keys();
println!("Creating directory {output_dir}");
std::fs::create_dir_all(output_dir).unwrap();
println!("Creating file {output_dir}/cks");
std::fs::write(format!("{output_dir}/cks"), client_key).unwrap();
std::fs::write(format!("{output_dir}/cks"), keys.client_key).unwrap();
println!("Creating file {output_dir}/pks");
std::fs::write(format!("{output_dir}/pks"), compact_key).unwrap();
std::fs::write(format!("{output_dir}/pks"), keys.compact_public_key).unwrap();
println!("Creating file {output_dir}/sks");
std::fs::write(format!("{output_dir}/sks"), server_key).unwrap();
std::fs::write(format!("{output_dir}/sks"), keys.server_key).unwrap();
}
25 changes: 14 additions & 11 deletions fhevm-engine/coprocessor/src/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
use crate::db_queries::{check_if_api_key_is_valid, check_if_ciphertexts_exist_in_db};
use crate::server::coprocessor::GenericResponse;
use crate::tfhe_ops::{
self, check_fhe_operand_types, current_ciphertext_version, debug_trivial_encrypt_be_bytes,
};
use fhevm_engine_common::tfhe_ops::{check_fhe_operand_types, current_ciphertext_version, debug_trivial_encrypt_be_bytes, deserialize_fhe_ciphertext};
use fhevm_engine_common::types::FhevmError;
use crate::types::CoprocessorError;
use crate::utils::sort_computations_by_dependencies;
use coprocessor::async_computation_input::Input;
Expand Down Expand Up @@ -167,7 +166,7 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
let mut decrypted: Vec<DebugDecryptResponseSingle> = Vec::with_capacity(cts.len());
for ct in cts {
let deserialized =
tfhe_ops::deserialize_fhe_ciphertext(ct.ciphertext_type, &ct.ciphertext)
deserialize_fhe_ciphertext(ct.ciphertext_type, &ct.ciphertext)
.unwrap();
decrypted.push(DebugDecryptResponseSingle {
output_type: ct.ciphertext_type as i32,
Expand Down Expand Up @@ -205,7 +204,7 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
let ciphertext_type: i16 = i_ct
.ciphertext_type
.try_into()
.map_err(|_e| CoprocessorError::UnknownFheType(i_ct.ciphertext_type))?;
.map_err(|_e| CoprocessorError::FhevmError(FhevmError::UnknownFheType(i_ct.ciphertext_type)))?;
let _ = sqlx::query!("
INSERT INTO ciphertexts(tenant_id, handle, ciphertext, ciphertext_version, ciphertext_type)
VALUES($1, $2, $3, $4, $5)
Expand Down Expand Up @@ -254,6 +253,7 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
let mut handle_types = Vec::with_capacity(comp.inputs.len());
let mut is_computation_scalar = false;
let mut this_comp_inputs: Vec<Vec<u8>> = Vec::with_capacity(comp.inputs.len());
let mut is_scalar_op_vec: Vec<bool> = Vec::with_capacity(comp.inputs.len());
for (idx, ih) in comp.inputs.iter().enumerate() {
if let Some(input) = &ih.input {
match input {
Expand All @@ -263,27 +263,30 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
.expect("this must be found if operand is non scalar");
handle_types.push(*ct_type);
this_comp_inputs.push(ih.clone());
is_scalar_op_vec.push(false);
}
Input::Scalar(sc) => {
is_computation_scalar = true;
handle_types.push(-1);
this_comp_inputs.push(sc.clone());
is_scalar_op_vec.push(true);
assert!(idx == 1, "we should have checked earlier that only second operand can be scalar");
}
}
}
}

computations_inputs.push(this_comp_inputs);
are_comps_scalar.push(is_computation_scalar);
// check before we insert computation that it has
// to succeed according to the type system
let output_type = check_fhe_operand_types(
comp.operation,
&handle_types,
is_computation_scalar,
&comp.inputs,
)?;
&this_comp_inputs,
&is_scalar_op_vec,
).map_err(|e| CoprocessorError::FhevmError(e))?;

computations_inputs.push(this_comp_inputs);
are_comps_scalar.push(is_computation_scalar);
// fill in types with output handles that are computed as we go
assert!(ct_types
.insert(comp.output_handle.clone(), output_type)
Expand All @@ -300,7 +303,7 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
let fhe_operation: i16 = comp
.operation
.try_into()
.map_err(|_| CoprocessorError::UnknownFheOperation(comp.operation))?;
.map_err(|_| CoprocessorError::FhevmError(FhevmError::UnknownFheOperation(comp.operation)))?;
let res = query!(
"
INSERT INTO computations(tenant_id, output_handle, dependencies, fhe_operation, is_completed, is_scalar)
Expand Down
8 changes: 3 additions & 5 deletions fhevm-engine/coprocessor/src/tests/operators.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,12 @@ use crate::tests::utils::wait_until_all_ciphertexts_computed;
use crate::{
server::coprocessor::{async_computation_input::Input, AsyncComputationInput},
tests::utils::{default_api_key, setup_test_app},
tfhe_ops::{
does_fhe_operation_support_both_encrypted_operands, does_fhe_operation_support_scalar,
},
types::{FheOperationType, SupportedFheOperations},
};
use bigdecimal::num_bigint::BigInt;
use std::{ops::Not, str::FromStr};
use fhevm_engine_common::tfhe_ops::{does_fhe_operation_support_both_encrypted_operands, does_fhe_operation_support_scalar};
use fhevm_engine_common::types::{FheOperationType, SupportedFheOperations};
use strum::IntoEnumIterator;
use std::{ops::Not, str::FromStr};
use tonic::metadata::MetadataValue;

struct BinaryOperatorTestCase {
Expand Down
15 changes: 11 additions & 4 deletions fhevm-engine/coprocessor/src/tfhe_worker.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
use crate::tfhe_ops::{
use fhevm_engine_common::tfhe_ops::{
current_ciphertext_version, deserialize_fhe_ciphertext, perform_fhe_operation,
};
use crate::types::{SupportedFheCiphertexts, TfheTenantKeys};
use crate::types::TfheTenantKeys;
use fhevm_engine_common::types::SupportedFheCiphertexts;
use sqlx::{postgres::PgListener, query, Acquire};
use std::{
cell::Cell,
Expand Down Expand Up @@ -223,13 +224,19 @@ async fn tfhe_worker_cycle(
} else {
deserialized_cts.push(
deserialize_fhe_ciphertext(*ct_type, ct_bytes.as_slice())
.map_err(|e| (e, w.tenant_id, w.output_handle.clone()))?,
.map_err(|e| {
let err: Box<dyn std::error::Error + Send + Sync> = Box::new(e);
(err, w.tenant_id, w.output_handle.clone())
})?,
);
}
}

let res = perform_fhe_operation(w.fhe_operation, &deserialized_cts)
.map_err(|e| (e, w.tenant_id, w.output_handle.clone()))?;
.map_err(|e| {
let err: Box<dyn std::error::Error + Send + Sync> = Box::new(e);
(err, w.tenant_id, w.output_handle.clone())
})?;
let (db_type, db_bytes) = res.serialize();

Ok((w, db_type, db_bytes))
Expand Down
Loading