feat(coprocessor): enable gpu backend in zkproof-worker (#1827)

Author: goshawk-3

.github/workflows/coprocessor-gpu-tests.yml

@@ -191,6 +191,7 @@ jobs:
           cargo test \
           -p tfhe-worker \
           -p sns-worker \
+          -p zkproof-worker \
           --release \
           --features=gpu \
           -- \

coprocessor/fhevm-engine/fhevm-engine-common/src/tenant_keys.rs

@@ -12,8 +12,13 @@ pub struct TfheTenantKeys {
     pub chain_id: i64,
     pub verifying_contract_address: String,
     pub acl_contract_address: String,
+
+    #[cfg(not(feature = "gpu"))]
     pub sks: tfhe::ServerKey,
 
+    #[cfg(feature = "gpu")]
+    pub sks: tfhe::CudaServerKey,
+
     pub pks: tfhe::CompactPublicKey,
     pub public_params: Arc<tfhe::zk::CompactPkeCrs>,
 }
@@ -23,7 +28,12 @@ pub struct FetchTenantKeyResult {
     pub chain_id: i64,
     pub verifying_contract_address: String,
     pub acl_contract_address: String,
+
+    #[cfg(not(feature = "gpu"))]
     pub server_key: tfhe::ServerKey,
+    #[cfg(feature = "gpu")]
+    pub server_key: tfhe::CudaServerKey,
+
     pub public_params: Arc<tfhe::zk::CompactPkeCrs>,
     pub pks: tfhe::CompactPublicKey,
 }
@@ -103,7 +113,7 @@ where
         #[cfg(feature = "gpu")]
         let sks = {
             let csks: tfhe::CompressedServerKey = safe_deserialize_key(&sks_key)?;
-            csks.decompress()
+            csks.decompress_to_gpu()
         };
         let pks: tfhe::CompactPublicKey = safe_deserialize_key(&pks_key)?;
         let public_params: tfhe::zk::CompactPkeCrs = safe_deserialize_key(&public_params_key)?;

coprocessor/fhevm-engine/fhevm-engine-common/src/tfhe_ops.rs

@@ -292,7 +292,8 @@ pub fn extract_ct_list(
     let mut res = Vec::new();
     for idx in 0..expanded.len() {
         let Some(data_kind) = expanded.get_kind_of(idx) else {
-            panic!("we're itering over what ciphertext told us how many ciphertexts are there, it must exist")
+            tracing::error!(len = expanded.len(), idx, "get_kind_of returned None");
+            continue;
         };
 
         match data_kind {

coprocessor/fhevm-engine/zkproof-worker/Cargo.toml

@@ -30,6 +30,8 @@ fhevm-engine-common = { path = "../fhevm-engine-common" }
 # crates.io dependencies
 [features]
 nightly-avx512 = ["tfhe/nightly-avx512"]
+gpu = ["tfhe/gpu", "fhevm-engine-common/gpu", "test-harness/gpu"]
+
 
 [dev-dependencies]
 serial_test = { workspace = true }

coprocessor/fhevm-engine/zkproof-worker/src/verifier.rs

@@ -125,7 +125,8 @@ pub async fn execute_verify_proofs_loop(
     conf: Config,
     last_active_at: Arc<RwLock<SystemTime>>,
 ) -> Result<(), ExecutionError> {
-    info!(conf = %conf, "Starting with config");
+    let gpu_enabled = fhevm_engine_common::utils::log_backend();
+    info!(gpu_enabled, conf = %conf, "Starting with config");
 
     // Tenants key cache is shared amongst all workers
     let tenant_key_cache = Arc::new(RwLock::new(LruCache::new(