test(coprocessor): add dependent whitepaper benchmark

Author: antoniupop

coprocessor/fhevm-engine/coprocessor/benches/erc20.rs

@@ -91,6 +91,19 @@ fn main() {
                 ));
             });
 
+            group.throughput(Throughput::Elements(num_elems));
+            let bench_id = format!(
+                "{bench_name}::throughput::dependent_whitepaper::FHEUint64::{num_elems}_elems::{bench_optimization_target}"
+            );
+            group.bench_with_input(bench_id.clone(), &num_elems, move |b, &num_elems| {
+                let _ = Runtime::new()
+                    .unwrap()
+                    .block_on(schedule_dependent_erc20_whitepaper(
+                        b,
+                        num_elems as usize,
+                        bench_id.clone(),
+                    ));
+            });
             group.throughput(Throughput::Elements(num_elems));
             let bench_id = format!(
                 "{bench_name}::throughput::dependent_no_cmux::FHEUint64::{num_elems}_elems::{bench_optimization_target}"
@@ -466,6 +479,208 @@ async fn schedule_erc20_no_cmux(
     Ok(())
 }
 
+async fn schedule_dependent_erc20_whitepaper(
+    bencher: &mut Bencher<'_, WallTime>,
+    num_tx: usize,
+    bench_id: String,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let app = setup_test_app().await?;
+    let pool = sqlx::postgres::PgPoolOptions::new()
+        .max_connections(2)
+        .connect(app.db_url())
+        .await?;
+    let mut client = FhevmCoprocessorClient::connect(app.app_url().to_string()).await?;
+
+    let mut handle_counter: u64 = random_handle();
+    let mut next_handle = || {
+        let out: u64 = handle_counter;
+        handle_counter += 1;
+        out.to_be_bytes().to_vec()
+    };
+    let api_key_header = format!("bearer {}", default_api_key());
+
+    let mut output_handles = vec![];
+    let mut async_computations = vec![];
+    let mut num_samples: usize = num_tx;
+    let samples = std::env::var("FHEVM_TEST_NUM_SAMPLES");
+    if let Ok(samples) = samples {
+        num_samples = samples.parse::<usize>().unwrap();
+    }
+
+    let keys = query_tenant_keys(vec![default_tenant_id()], &pool)
+        .await
+        .map_err(|e| {
+            let e: Box<dyn std::error::Error> = e;
+            e
+        })?;
+    let keys = &keys[0];
+
+    let mut builder = tfhe::ProvenCompactCiphertextList::builder(&keys.pks);
+    let the_list = builder
+        .push(20_u64) // Initial balance destination
+        .build_with_proof_packed(&keys.public_params, &[], tfhe::zk::ZkComputeLoad::Proof)
+        .unwrap();
+    let serialized = safe_serialize(&the_list);
+    let mut input_request = tonic::Request::new(InputUploadBatch {
+        input_ciphertexts: vec![InputToUpload {
+            input_payload: serialized,
+            signatures: Vec::new(),
+            user_address: test_random_user_address(),
+            contract_address: test_random_contract_address(),
+        }],
+    });
+    input_request.metadata_mut().append(
+        "authorization",
+        MetadataValue::from_str(&api_key_header).unwrap(),
+    );
+    let resp = client.upload_inputs(input_request).await?;
+    let resp = resp.get_ref();
+    assert_eq!(resp.upload_responses.len(), 1);
+    let first_resp = &resp.upload_responses[0];
+    assert_eq!(first_resp.input_handles.len(), 1);
+    let handle_bald = first_resp.input_handles[0].handle.clone();
+    let mut bald = AsyncComputationInput {
+        input: Some(Input::InputHandle(handle_bald.clone())),
+    };
+
+    let mut builder = tfhe::ProvenCompactCiphertextList::builder(&keys.pks);
+    let the_list = builder
+        .push(100_u64) // Balance source
+        .push(10_u64) // Transfer amount
+        .build_with_proof_packed(&keys.public_params, &[], tfhe::zk::ZkComputeLoad::Proof)
+        .unwrap();
+
+    let serialized = safe_serialize(&the_list);
+    let mut input_request = tonic::Request::new(InputUploadBatch {
+        input_ciphertexts: vec![InputToUpload {
+            input_payload: serialized,
+            signatures: Vec::new(),
+            user_address: test_random_user_address(),
+            contract_address: test_random_contract_address(),
+        }],
+    });
+    input_request.metadata_mut().append(
+        "authorization",
+        MetadataValue::from_str(&api_key_header).unwrap(),
+    );
+    let resp = client.upload_inputs(input_request).await?;
+    let resp = resp.get_ref();
+    assert_eq!(resp.upload_responses.len(), 1);
+    let first_resp = &resp.upload_responses[0];
+    assert_eq!(first_resp.input_handles.len(), 2);
+
+    for _ in 0..=(num_samples - 1) as u32 {
+        let handle_bals = first_resp.input_handles[0].handle.clone();
+        let bals = AsyncComputationInput {
+            input: Some(Input::InputHandle(handle_bals.clone())),
+        };
+        let handle_trxa = first_resp.input_handles[1].handle.clone();
+        let trxa = AsyncComputationInput {
+            input: Some(Input::InputHandle(handle_trxa.clone())),
+        };
+
+        let has_enough_funds_handle = next_handle();
+        output_handles.push(has_enough_funds_handle.clone());
+        let new_to_amount_target_handle = next_handle();
+        output_handles.push(new_to_amount_target_handle.clone());
+        let new_to_amount_handle = next_handle();
+        output_handles.push(new_to_amount_handle.clone());
+        let new_from_amount_target_handle = next_handle();
+        output_handles.push(new_from_amount_target_handle.clone());
+        let new_from_amount_handle = next_handle();
+        output_handles.push(new_from_amount_handle.clone());
+
+        async_computations.push(AsyncComputation {
+            operation: FheOperation::FheGe.into(),
+            output_handle: has_enough_funds_handle.clone(),
+            inputs: vec![bals.clone(), trxa.clone()],
+        });
+        async_computations.push(AsyncComputation {
+            operation: FheOperation::FheAdd.into(),
+            output_handle: new_to_amount_target_handle.clone(),
+            inputs: vec![bald.clone(), trxa.clone()],
+        });
+        async_computations.push(AsyncComputation {
+            operation: FheOperation::FheIfThenElse.into(),
+            output_handle: new_to_amount_handle.clone(),
+            inputs: vec![
+                AsyncComputationInput {
+                    input: Some(Input::InputHandle(has_enough_funds_handle.clone())),
+                },
+                AsyncComputationInput {
+                    input: Some(Input::InputHandle(new_to_amount_target_handle.clone())),
+                },
+                bald.clone(),
+            ],
+        });
+        async_computations.push(AsyncComputation {
+            operation: FheOperation::FheSub.into(),
+            output_handle: new_from_amount_target_handle.clone(),
+            inputs: vec![bals.clone(), trxa.clone()],
+        });
+        async_computations.push(AsyncComputation {
+            operation: FheOperation::FheIfThenElse.into(),
+            output_handle: new_from_amount_handle.clone(),
+            inputs: vec![
+                AsyncComputationInput {
+                    input: Some(Input::InputHandle(has_enough_funds_handle.clone())),
+                },
+                AsyncComputationInput {
+                    input: Some(Input::InputHandle(new_from_amount_target_handle.clone())),
+                },
+                bals.clone(),
+            ],
+        });
+
+        bald = AsyncComputationInput {
+            input: Some(Input::InputHandle(new_to_amount_handle.clone())),
+        };
+    }
+
+    let mut compute_request = tonic::Request::new(AsyncComputeRequest {
+        computations: async_computations,
+    });
+    compute_request.metadata_mut().append(
+        "authorization",
+        MetadataValue::from_str(&api_key_header).unwrap(),
+    );
+    let _resp = client.clone().async_compute(compute_request).await.unwrap();
+    let app_ref = &app;
+    bencher
+        .to_async(FuturesExecutor)
+        .iter_custom(|iters| async move {
+            let db_url = app_ref.db_url().to_string();
+            let now = SystemTime::now();
+            let _ = tokio::task::spawn_blocking(move || {
+                Runtime::new()
+                    .unwrap()
+                    .block_on(async { wait_until_all_ciphertexts_computed(db_url).await.unwrap() });
+                println!(
+                    "Execution time: {} -- {}",
+                    now.elapsed().unwrap().as_millis(),
+                    TIMING.load(std::sync::atomic::Ordering::SeqCst) / 1000
+                );
+            })
+            .await;
+            std::time::Duration::from_micros(
+                TIMING.swap(0, std::sync::atomic::Ordering::SeqCst) * iters,
+            )
+        });
+
+    let params = keys.cks.computation_parameters();
+    write_to_json::<u64, _>(
+        &bench_id,
+        params,
+        "",
+        "erc20-transfer",
+        &OperatorType::Atomic,
+        64,
+        vec![],
+    );
+
+    Ok(())
+}
+
 async fn schedule_dependent_erc20_no_cmux(
     bencher: &mut Bencher<'_, WallTime>,
     num_tx: usize,