Fix support for 512bit + inputs

zama-ai · Sep 9, 2024 · 30c57e3 · 30c57e3
1 parent ad917d1
commit 30c57e3
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 28 deletions.
diff --git a/fhevm-engine/Cargo.lock b/fhevm-engine/Cargo.lock
diff --git a/fhevm-engine/coprocessor/Cargo.toml b/fhevm-engine/coprocessor/Cargo.toml
@@ -4,9 +4,9 @@ version = "0.1.0"
 edition = "2021"
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
-tfhe = { version = "0.8.0-alpha.2", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
+tfhe = { version = "0.8.0-alpha.4", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
 [target.'cfg(target_arch = "aarch64")'.dependencies]
-tfhe = { version = "0.8.0-alpha.2", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
+tfhe = { version = "0.8.0-alpha.4", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
 
 [dependencies]
 # Common dependencies

diff --git a/fhevm-engine/coprocessor/src/tests/inputs.rs b/fhevm-engine/coprocessor/src/tests/inputs.rs
@@ -1,6 +1,6 @@
 use std::str::FromStr;
 
-use tfhe::integer::U256;
+use tfhe::integer::{bigint::StaticUnsignedBigInt, U256};
 use tonic::metadata::MetadataValue;
 
 use crate::{db_queries::query_tenant_keys, server::coprocessor::{fhevm_coprocessor_client::FhevmCoprocessorClient, DebugDecryptRequest, InputToUpload, InputUploadBatch}, tests::utils::{default_api_key, default_tenant_id, setup_test_app}};
@@ -32,9 +32,9 @@ async fn test_fhe_inputs() -> Result<(), Box<dyn std::error::Error>> {
         .push(5u128)
         // TODO: 160 bits test
         .push(U256::from(7u32))
-        // .push(StaticUnsignedBigInt::<8>::from(8u32))
-        // .push(StaticUnsignedBigInt::<16>::from(9u32))
-        // .push(StaticUnsignedBigInt::<32>::from(10u32))
+        .push(StaticUnsignedBigInt::<8>::from(8u32))
+        .push(StaticUnsignedBigInt::<16>::from(9u32))
+        .push(StaticUnsignedBigInt::<32>::from(10u32))
         .build();
 
     let serialized = bincode::serialize(&the_list).unwrap();
@@ -58,7 +58,7 @@ async fn test_fhe_inputs() -> Result<(), Box<dyn std::error::Error>> {
 
     let first_resp = &resp.upload_responses[0];
 
-    assert_eq!(first_resp.input_handles.len(), 7);
+    assert_eq!(first_resp.input_handles.len(), 10);
 
     let mut decr_handles: Vec<Vec<u8>> = Vec::new();
     for handle in &first_resp.input_handles {
@@ -74,9 +74,7 @@ async fn test_fhe_inputs() -> Result<(), Box<dyn std::error::Error>> {
     );
     let resp = client.debug_decrypt_ciphertext(decrypt_request).await?;
     let resp = resp.get_ref();
-    assert_eq!(resp.values.len(), 7);
-    // TODO: add ebytes support, we only support 256 bits max
-    // assert_eq!(resp.values.len(), 10);
+    assert_eq!(resp.values.len(), 10);
 
     assert_eq!(resp.values[0].output_type, 0);
     assert_eq!(resp.values[0].value, "false");
@@ -92,12 +90,12 @@ async fn test_fhe_inputs() -> Result<(), Box<dyn std::error::Error>> {
     assert_eq!(resp.values[5].value, "5");
     assert_eq!(resp.values[6].output_type, 8);
     assert_eq!(resp.values[6].value, "7");
-    // assert_eq!(resp.values[7].output_type, 9);
-    // assert_eq!(resp.values[7].value, "8");
-    // assert_eq!(resp.values[8].output_type, 10);
-    // assert_eq!(resp.values[8].value, "9");
-    // assert_eq!(resp.values[9].output_type, 11);
-    // assert_eq!(resp.values[9].value, "10");
+    assert_eq!(resp.values[7].output_type, 9);
+    assert_eq!(resp.values[7].value, "8");
+    assert_eq!(resp.values[8].output_type, 10);
+    assert_eq!(resp.values[8].value, "9");
+    assert_eq!(resp.values[9].output_type, 11);
+    assert_eq!(resp.values[9].value, "10");
 
     Ok(())
 }

diff --git a/fhevm-engine/executor/Cargo.toml b/fhevm-engine/executor/Cargo.toml
@@ -4,9 +4,9 @@ version = "0.1.0"
 edition = "2021"
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
-tfhe = { version = "0.8.0-alpha.2", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
+tfhe = { version = "0.8.0-alpha.4", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
 [target.'cfg(target_arch = "aarch64")'.dependencies]
-tfhe = { version = "0.8.0-alpha.2", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
+tfhe = { version = "0.8.0-alpha.4", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
 
 [dependencies]
 clap.workspace = true

diff --git a/fhevm-engine/fhevm-engine-common/Cargo.toml b/fhevm-engine/fhevm-engine-common/Cargo.toml
@@ -4,9 +4,9 @@ version = "0.1.0"
 edition = "2021"
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
-tfhe = { version = "0.8.0-alpha.2", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
+tfhe = { version = "0.8.0-alpha.4", features = ["boolean", "shortint", "integer", "x86_64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
 [target.'cfg(target_arch = "aarch64")'.dependencies]
-tfhe = { version = "0.8.0-alpha.2", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
+tfhe = { version = "0.8.0-alpha.4", features = ["boolean", "shortint", "integer", "aarch64-unix", "zk-pok", "experimental-force_fft_algo_dif4"] }
 
 [dependencies]
 sha3.workspace = true

diff --git a/fhevm-engine/fhevm-engine-common/src/tfhe_ops.rs b/fhevm-engine/fhevm-engine-common/src/tfhe_ops.rs
@@ -272,6 +272,30 @@ pub fn try_expand_ciphertext_list(
 
                 res.push(SupportedFheCiphertexts::FheUint256(ct));
             }
+            tfhe::FheTypes::Uint512 => {
+                let ct: tfhe::FheUint512 = expanded
+                    .get(idx)
+                    .expect("Index must exist")
+                    .expect("Must succeed, we just checked this is the type");
+
+                res.push(SupportedFheCiphertexts::FheBytes64(ct));
+            }
+            tfhe::FheTypes::Uint1024 => {
+                let ct: tfhe::FheUint1024 = expanded
+                    .get(idx)
+                    .expect("Index must exist")
+                    .expect("Must succeed, we just checked this is the type");
+
+                res.push(SupportedFheCiphertexts::FheBytes128(ct));
+            }
+            tfhe::FheTypes::Uint2048 => {
+                let ct: tfhe::FheUint2048 = expanded
+                    .get(idx)
+                    .expect("Index must exist")
+                    .expect("Must succeed, we just checked this is the type");
+
+                res.push(SupportedFheCiphertexts::FheBytes256(ct));
+            }
             other => {
                 return Err(FhevmError::CiphertextExpansionUnsupportedCiphertextKind(
                     other,