chore(zk): add benches to tfhe-zk-pok

.github/workflows/benchmark_tfhe_zk_pok.yml

@@ -0,0 +1,173 @@
+# Run benchmarks of the tfhe-zk-pok crate on an instance and return parsed results to Slab CI bot.
+name: tfhe-zk-pok benchmarks
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+  schedule:
+    # Weekly benchmarks will be triggered each Saturday at 3a.m.
+    - cron: '0 3 * * 6'
+env:
+  CARGO_TERM_COLOR: always
+  RESULTS_FILENAME: parsed_benchmark_results_${{ github.sha }}.json
+  PARSE_INTEGER_BENCH_CSV_FILE: tfhe_rs_integer_benches_${{ github.sha }}.csv
+  ACTION_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
+  RUST_BACKTRACE: "full"
+  RUST_MIN_STACK: "8388608"
+  SLACK_CHANNEL: ${{ secrets.SLACK_CHANNEL }}
+  SLACK_ICON: https://pbs.twimg.com/profile_images/1274014582265298945/OjBKP9kn_400x400.png
+  SLACK_USERNAME: ${{ secrets.BOT_USERNAME }}
+  SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }}
+
+jobs:
+  should-run:
+    runs-on: ubuntu-latest
+    if: github.event_name == 'workflow_dispatch' ||
+      ((github.event_name == 'push' || github.event_name == 'schedule') && github.repository == 'zama-ai/tfhe-rs')
+    outputs:
+      zk_pok_changed: ${{ steps.changed-files.outputs.zk_pok_any_changed }}
+    steps:
+      - name: Checkout tfhe-rs
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
+        with:
+          fetch-depth: 0
+
+      - name: Check for file changes
+        id: changed-files
+        uses: tj-actions/changed-files@c3a1bb2c992d77180ae65be6ae6c166cf40f857c
+        with:
+          since_last_remote_commit: true
+          files_yaml: |
+            zk_pok:
+              - tfhe-zk-pok/**
+              - .github/workflows/benchmark_tfhe_zk_pok.yml
+
+  setup-instance:
+    name: Setup instance (tfhe-zk-pok-benchmarks)
+    runs-on: ubuntu-latest
+    needs: should-run
+    if: github.event_name == 'workflow_dispatch' ||
+      (github.event_name == 'schedule' && github.repository == 'zama-ai/tfhe-rs') ||
+      (github.event_name == 'push' &&
+      github.repository == 'zama-ai/tfhe-rs' &&
+      needs.should-run.outputs.zk_pok_changed == 'true')
+    outputs:
+      runner-name: ${{ steps.start-instance.outputs.label }}
+    steps:
+      - name: Start instance
+        id: start-instance
+        uses: zama-ai/slab-github-runner@801df0b8db5ea2b06128b7476c652f5ed5f193a8
+        with:
+          mode: start
+          github-token: ${{ secrets.SLAB_ACTION_TOKEN }}
+          slab-url: ${{ secrets.SLAB_BASE_URL }}
+          job-secret: ${{ secrets.JOB_SECRET }}
+          backend: aws
+          profile: bench
+
+  tfhe-zk-pok-benchmarks:
+    name: Execute tfhe-zk-pok benchmarks
+    if: needs.setup-instance.result != 'skipped'
+    needs: setup-instance
+    concurrency:
+      group: ${{ github.workflow }}_${{github.event_name}}_${{ github.ref }}${{ github.ref == 'refs/heads/main' && github.sha || '' }}
+      cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+    runs-on: ${{ needs.setup-instance.outputs.runner-name }}
+    steps:
+      - name: Checkout tfhe-rs repo with tags
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
+        with:
+          fetch-depth: 0
+          token: ${{ secrets.FHE_ACTIONS_TOKEN }}
+
+      - name: Get benchmark details
+        run: |
+          {
+            echo "BENCH_DATE=$(date --iso-8601=seconds)";
+            echo "COMMIT_DATE=$(git --no-pager show -s --format=%cd --date=iso8601-strict ${{ github.sha }})";
+            echo "COMMIT_HASH=$(git describe --tags --dirty)";
+          } >> "${GITHUB_ENV}"
+
+      - name: Install rust
+        uses: dtolnay/rust-toolchain@7b1c307e0dcbda6122208f10795a713336a9b35a
+        with:
+          toolchain: nightly
+
+      - name: Checkout Slab repo
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
+        with:
+          repository: zama-ai/slab
+          path: slab
+          token: ${{ secrets.FHE_ACTIONS_TOKEN }}
+
+      - name: Run benchmarks
+        run: |
+          make bench_tfhe_zk_pok
+
+      - name: Parse results
+        run: |
+          python3 ./ci/benchmark_parser.py target/criterion ${{ env.RESULTS_FILENAME }} \
+          --database tfhe_rs \
+          --crate tfhe-zk-pok \
+          --hardware "hpc7a.96xlarge" \
+          --backend cpu \
+          --project-version "${{ env.COMMIT_HASH }}" \
+          --branch ${{ github.ref_name }} \
+          --commit-date "${{ env.COMMIT_DATE }}" \
+          --bench-date "${{ env.BENCH_DATE }}" \
+          --walk-subdirs \
+          --name-suffix avx512 \
+          --throughput
+
+      - name: Upload parsed results artifact
+        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882
+        with:
+          name: ${{ github.sha }}_tfhe_zk_pok
+          path: ${{ env.RESULTS_FILENAME }}
+
+      - name: Checkout Slab repo
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
+        with:
+          repository: zama-ai/slab
+          path: slab
+          token: ${{ secrets.FHE_ACTIONS_TOKEN }}
+
+      - name: Send data to Slab
+        shell: bash
+        run: |
+          python3 slab/scripts/data_sender.py ${{ env.RESULTS_FILENAME }} "${{ secrets.JOB_SECRET }}" \
+          --slab-url "${{ secrets.SLAB_URL }}"
+
+      - name: Slack Notification
+        if: ${{ failure() }}
+        continue-on-error: true
+        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
+        env:
+          SLACK_COLOR: ${{ job.status }}
+          SLACK_MESSAGE: "tfhe-zk-pok benchmarks finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"
+
+  teardown-instance:
+    name: Teardown instance (tfhe-zk-pok-benchmarks)
+    if: ${{ always() && needs.setup-instance.result != 'skipped' }}
+    needs: [ setup-instance, tfhe-zk-pok-benchmarks ]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Stop instance
+        id: stop-instance
+        uses: zama-ai/slab-github-runner@801df0b8db5ea2b06128b7476c652f5ed5f193a8
+        with:
+          mode: stop
+          github-token: ${{ secrets.SLAB_ACTION_TOKEN }}
+          slab-url: ${{ secrets.SLAB_BASE_URL }}
+          job-secret: ${{ secrets.JOB_SECRET }}
+          label: ${{ needs.setup-instance.outputs.runner-name }}
+
+      - name: Slack Notification
+        if: ${{ failure() }}
+        continue-on-error: true
+        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
+        env:
+          SLACK_COLOR: ${{ job.status }}
+          SLACK_MESSAGE: "Instance teardown (tfhe-zk-pok-benchmarks) finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"

.gitignore

@@ -13,6 +13,7 @@ target/
 
 # Some of our bench outputs
 /tfhe/benchmarks_parameters
+/tfhe-zk-pok/benchmarks_parameters
 **/*.csv
 
 # dieharder run log

Makefile

@@ -1191,6 +1191,11 @@ bench_hlapi_erc20_gpu: install_rs_check_toolchain
 	--bench hlapi-erc20 \
 	--features=$(TARGET_ARCH_FEATURE),integer,gpu,internal-keycache,pbs-stats,nightly-avx512 -p $(TFHE_SPEC) --
 
+.PHONY: bench_tfhe_zk_pok # Run benchmarks for the tfhe_zk_pok crate
+bench_tfhe_zk_pok: install_rs_check_toolchain
+	RUSTFLAGS="$(RUSTFLAGS)" \
+	cargo $(CARGO_RS_CHECK_TOOLCHAIN) bench -p tfhe-zk-pok --
+
 #
 # Utility tools
 #

ci/benchmark_parser.py

@@ -94,10 +94,17 @@
     default="cpu",
     help="Backend on which benchmarks have run",
 )
+parser.add_argument(
+    "--crate",
+    dest="crate",
+    default="tfhe",
+    help="Crate for which benchmarks have run",
+)
 
 
 def recursive_parse(
     directory,
+    crate,
     walk_subdirs=False,
     name_suffix="",
     compute_throughput=False,
@@ -108,6 +115,7 @@ def recursive_parse(
     .json extension at the top-level of this directory.
 
     :param directory: path to directory that contains raw results as :class:`pathlib.Path`
+    :param crate: the name of the crate that has been benched
     :param walk_subdirs: traverse results subdirectories if parameters changes for benchmark case.
     :param name_suffix: a :class:`str` suffix to apply to each test name found
     :param compute_throughput: compute number of operations per second and operations per
@@ -143,7 +151,7 @@ def recursive_parse(
                 continue
 
             try:
-                params, display_name, operator = get_parameters(test_name)
+                params, display_name, operator = get_parameters(test_name, crate)
             except Exception as err:
                 parsing_failures.append((full_name, f"failed to get parameters: {err}"))
                 continue
@@ -276,7 +284,7 @@ def _parse_key_results(result_file, bench_type):
         reader = csv.reader(csv_file)
         for test_name, value in reader:
             try:
-                params, display_name, operator = get_parameters(test_name)
+                params, display_name, operator = get_parameters(test_name, crate)
             except Exception as err:
                 parsing_failures.append((test_name, f"failed to get parameters: {err}"))
                 continue
@@ -318,15 +326,16 @@ def parse_key_gen_time(result_file):
     return _parse_key_results(result_file, "latency")
 
 
-def get_parameters(bench_id):
+def get_parameters(bench_id, directory):
     """
     Get benchmarks parameters recorded for a given benchmark case.
 
     :param bench_id: function name used for the benchmark case
+    :param directory: directory where the parameters are stored
 
     :return: :class:`tuple` as ``(benchmark parameters, display name, operator type)``
     """
-    params_dir = pathlib.Path("tfhe", "benchmarks_parameters", bench_id)
+    params_dir = pathlib.Path(directory, "benchmarks_parameters", bench_id)
     params = _parse_file_to_json(params_dir, "parameters.json")
 
     display_name = params.pop("display_name")
@@ -459,6 +468,7 @@ def check_mandatory_args(input_args):
 
         results, failures = recursive_parse(
             raw_results,
+            args.crate,
             args.walk_subdirs,
             args.name_suffix,
             args.throughput,

tfhe-zk-pok/Cargo.toml

@@ -28,3 +28,12 @@ tfhe-versionable = { version = "0.3.2", path = "../utils/tfhe-versionable" }
 serde_json = "~1.0"
 itertools = "0.11.0"
 bincode = "1.3.3"
+criterion = "0.5.1"
+
+[[bench]]
+name = "pke_v1"
+harness = false
+
+[[bench]]
+name = "pke_v2"
+harness = false

tfhe-zk-pok/benches/pke_v1.rs

@@ -0,0 +1,90 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use tfhe_zk_pok::proofs::pke::{prove, verify};
+use tfhe_zk_pok::proofs::ComputeLoad;
+use utils::{write_to_json, PKEV1_TEST_PARAMS, PKEV2_TEST_PARAMS};
+
+#[path = "./utils.rs"]
+mod utils;
+
+use crate::utils::init_params_v1;
+
+fn bench_pke_v1_prove(c: &mut Criterion) {
+    let bench_shortname = "pke_zk_proof_v1";
+    let bench_name = format!("tfhe_zk_pok::{bench_shortname}");
+    let mut bench_group = c.benchmark_group(&bench_name);
+    bench_group
+        .sample_size(15)
+        .measurement_time(std::time::Duration::from_secs(60));
+
+    let rng = &mut rand::thread_rng();
+
+    for (params, param_name) in [
+        (PKEV1_TEST_PARAMS, "PKEV1_TEST_PARAMS"),
+        (PKEV2_TEST_PARAMS, "PKEV2_TEST_PARAMS"),
+    ] {
+        let (public_param, public_commit, private_commit, metadata) = init_params_v1(params);
+        let effective_t = params.t >> 1;
+        let bits = (params.k as u32) * effective_t.ilog2();
+
+        for load in [ComputeLoad::Proof, ComputeLoad::Verify] {
+            let bench_id = format!("{bench_name}::{param_name}_{bits}_bits_packed_{load}");
+
+            bench_group.bench_function(&bench_id, |b| {
+                b.iter(|| {
+                    prove(
+                        (&public_param, &public_commit),
+                        &private_commit,
+                        &metadata,
+                        load,
+                        rng,
+                    )
+                })
+            });
+
+            write_to_json(&bench_id, params, param_name, bench_shortname);
+        }
+    }
+}
+
+fn bench_pke_v1_verify(c: &mut Criterion) {
+    let bench_shortname = "pke_zk_verify_v1";
+    let bench_name = format!("tfhe_zk_pok::{bench_shortname}");
+    let mut bench_group = c.benchmark_group(&bench_name);
+    bench_group
+        .sample_size(15)
+        .measurement_time(std::time::Duration::from_secs(60));
+
+    let rng = &mut rand::thread_rng();
+
+    for (params, param_name) in [
+        (PKEV1_TEST_PARAMS, "PKEV1_TEST_PARAMS"),
+        (PKEV2_TEST_PARAMS, "PKEV2_TEST_PARAMS"),
+    ] {
+        let (public_param, public_commit, private_commit, metadata) = init_params_v1(params);
+        let effective_t = params.t >> 1;
+        let bits = (params.k as u32) * effective_t.ilog2();
+
+        for load in [ComputeLoad::Proof, ComputeLoad::Verify] {
+            let bench_id = format!("{bench_name}::{param_name}_{bits}_bits_packed_{load}");
+
+            let proof = prove(
+                (&public_param, &public_commit),
+                &private_commit,
+                &metadata,
+                load,
+                rng,
+            );
+
+            bench_group.bench_function(&bench_id, |b| {
+                b.iter(|| {
+                    verify(&proof, (&public_param, &public_commit), &metadata).unwrap();
+                })
+            });
+
+            write_to_json(&bench_id, params, param_name, bench_shortname);
+        }
+    }
+}
+
+criterion_group!(benches_pke_v1, bench_pke_v1_verify, bench_pke_v1_prove);
+criterion_main!(benches_pke_v1);