chore(ci): measure and report key sizes used in benchmarks

Size of boostrapping and key switching keys used in benchmarks are
measured and then sent to Slab to be stored into our benchmark
database.

.github/workflows/boolean_benchmark.yml

@@ -64,7 +64,7 @@ jobs:
 
       - name: Run benchmarks
         run: |
-          RUSTFLAGS="-C target-cpu=native" cargo +nightly bench --bench boolean-bench --features=x86_64-unix,boolean,internal-keycache -p tfhe
+          make bench_boolean
 
       - name: Parse results
         run: |
@@ -84,7 +84,7 @@ jobs:
 
       - name: Run benchmarks with AVX512
         run: |
-          RUSTFLAGS="-C target-cpu=native" cargo +nightly bench --bench boolean-bench --features=x86_64-unix,boolean,internal-keycache,nightly-avx512 -p tfhe
+          make AVX512_SUPPORT=ON bench_boolean
 
       - name: Parse AVX512 results
         run: |
@@ -92,6 +92,16 @@ jobs:
           --name-suffix avx512 \
           --append-results
 
+      - name: Measure key sizes
+        run: |
+          make measure_boolean_key_sizes
+
+      - name: Parse key sizes results
+        run: |
+          python3 ./ci/benchmark_parser.py tfhe/boolean_key_sizes.csv ${{ env.RESULTS_FILENAME }} \
+          --key-sizes \
+          --append-results
+  
       - name: Upload parsed results artifact
         uses: actions/upload-artifact@v3
         with:

.github/workflows/shortint_benchmark.yml

@@ -64,7 +64,7 @@ jobs:
 
       - name: Run benchmarks
         run: |
-          RUSTFLAGS="-C target-cpu=native" cargo +nightly bench --bench shortint-bench --features=x86_64-unix,shortint,internal-keycache -p tfhe
+          make bench_shortint
 
       - name: Parse results
         run: |
@@ -85,7 +85,7 @@ jobs:
 
       - name: Run benchmarks with AVX512
         run: |
-          RUSTFLAGS="-C target-cpu=native" cargo +nightly bench --bench shortint-bench --features=x86_64-unix,shortint,internal-keycache,nightly-avx512 -p tfhe
+          make AVX512_SUPPORT=ON bench_shortint
 
       - name: Parse AVX512 results
         run: |
@@ -94,6 +94,16 @@ jobs:
           --name-suffix avx512 \
           --append-results
 
+      - name: Measure key sizes
+        run: |
+          make measure_shortint_key_sizes
+
+      - name: Parse key sizes results
+        run: |
+          python3 ./ci/benchmark_parser.py tfhe/shortint_key_sizes.csv ${{ env.RESULTS_FILENAME }} \
+          --key-sizes \
+          --append-results
+
       - name: Upload parsed results artifact
         uses: actions/upload-artifact@v3
         with:

Makefile

@@ -6,11 +6,18 @@ RS_BUILD_TOOLCHAIN:=$(shell \
 	( (echo $(TARGET_ARCH_FEATURE) | grep -q x86) && echo stable) || echo $(RS_CHECK_TOOLCHAIN))
 CARGO_RS_BUILD_TOOLCHAIN:=+$(RS_BUILD_TOOLCHAIN)
 MIN_RUST_VERSION:=1.65
+AVX512_SUPPORT?=OFF
 WASM_RUSTFLAGS:=
 # This is done to avoid forgetting it, we still precise the RUSTFLAGS in the commands to be able to
 # copy paste the command in the terminal and change them if required without forgetting the flags
 export RUSTFLAGS:=-C target-cpu=native
 
+ifeq ($(AVX512_SUPPORT),ON)
+		AVX512_FEATURE=nightly-avx512
+else
+		AVX512_FEATURE=
+endif
+
 .PHONY: rs_check_toolchain # Echo the rust toolchain used for checks
 rs_check_toolchain:
 	@echo $(RS_CHECK_TOOLCHAIN)
@@ -215,6 +222,30 @@ no_tfhe_typo:
 		exit 1; \
 	fi
 
+.PHONY: bench_shortint # Run benchmarks for shortint
+bench_shortint: install_rs_check_toolchain
+	RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_CHECK_TOOLCHAIN) bench \
+	--bench shortint-bench \
+	--features=$(TARGET_ARCH_FEATURE),shortint,internal-keycache,$(AVX512_FEATURE) -p tfhe
+
+.PHONY: bench_boolean # Run benchmarks for boolean
+bench_boolean: install_rs_check_toolchain
+	RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_CHECK_TOOLCHAIN) bench \
+	--bench boolean-bench \
+	--features=$(TARGET_ARCH_FEATURE),boolean,internal-keycache,$(AVX512_FEATURE) -p tfhe
+
+.PHONY: measure_shortint_key_sizes # Measure sizes of bootstrapping and key switching keys for shortint
+measure_shortint_key_sizes: install_rs_check_toolchain
+	RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_CHECK_TOOLCHAIN) run \
+	--example shortint_key_sizes \
+	--features=$(TARGET_ARCH_FEATURE),shortint,internal-keycache
+
+.PHONY: measure_boolean_key_sizes # Measure sizes of bootstrapping and key switching keys for boolean
+measure_boolean_key_sizes: install_rs_check_toolchain
+	RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_CHECK_TOOLCHAIN) run \
+	--example boolean_key_sizes \
+	--features=$(TARGET_ARCH_FEATURE),boolean,internal-keycache
+
 .PHONY: pcc # pcc stands for pre commit checks
 pcc: no_tfhe_typo check_fmt doc clippy_all check_compile_tests
 

ci/benchmark_parser.py

@@ -2,17 +2,20 @@
 benchmark_parser
 ----------------
 
-Parse criterion benchmark results.
+Parse criterion benchmark or keys size results.
 """
 import argparse
+import csv
 import pathlib
 import json
 import sys
 
 
 parser = argparse.ArgumentParser()
-parser.add_argument('results_dir',
-                    help='Location of criterion benchmark results directory')
+parser.add_argument('results',
+                    help='Location of criterion benchmark results directory.'
+                         'If the --key-size option is used, then the value would have to point to'
+                         'a CSV file.')
 parser.add_argument('output_file', help='File storing parsed results')
 parser.add_argument('-d', '--database', dest='database',
                     help='Name of the database used to store results')
@@ -32,6 +35,8 @@ parser.add_argument('--append-results', dest='append_results', action='store_tru
                     help='Append parsed results to an existing file')
 parser.add_argument('--walk-subdirs', dest='walk_subdirs', action='store_true',
                     help='Check for results in subdirectories')
+parser.add_argument('--key-sizes', dest='key_sizes', action='store_true',
+                    help='Parse only the results regarding keys size measurments')
 
 
 def recursive_parse(directory, walk_subdirs=False, name_suffix=""):
@@ -74,8 +79,8 @@ def parse_benchmark_file(directory):
 
     :return: name of the test as :class:`str`
     """
-    raw_results = _parse_file_to_json(directory, "benchmark.json")
-    return raw_results["full_id"].replace(" ", "_")
+    raw_res = _parse_file_to_json(directory, "benchmark.json")
+    return raw_res["full_id"].replace(" ", "_")
 
 
 def parse_estimate_file(directory):
@@ -86,13 +91,30 @@ def parse_estimate_file(directory):
 
     :return: :class:`dict` of data points
     """
-    raw_results = _parse_file_to_json(directory, "estimates.json")
+    raw_res = _parse_file_to_json(directory, "estimates.json")
     return {
-        stat_name: raw_results[stat_name]["point_estimate"]
+        stat_name: raw_res[stat_name]["point_estimate"]
         for stat_name in ("mean", "std_dev")
     }
 
 
+def parse_key_sizes(result_file):
+    """
+    Parse file containing key sizes results. The file must be formatted as CSV.
+
+    :param result_file: results file as :class:`pathlib.Path`
+
+    :return: :class:`list` of data points
+    """
+    result_values = list()
+    with result_file.open() as csv_file:
+        reader = csv.reader(csv_file)
+        for (test_name, value) in reader:
+            result_values.append({"value": int(value), "test": test_name})
+
+    return result_values
+
+
 def _parse_file_to_json(directory, filename):
     result_file = directory.joinpath(filename)
     return json.loads(result_file.read_text())
@@ -137,7 +159,7 @@ def check_mandatory_args(input_args):
     missing_args = list()
     for arg_name in vars(input_args):
         if arg_name in ["results_dir", "output_file", "name_suffix",
-                        "append_results", "walk_subdirs"]:
+                        "append_results", "walk_subdirs", "key_sizes"]:
             continue
         if not getattr(input_args, arg_name):
             missing_args.append(arg_name)
@@ -152,8 +174,13 @@ if __name__ == "__main__":
     args = parser.parse_args()
     check_mandatory_args(args)
 
-    print("Parsing benchmark results... ")
-    results = recursive_parse(pathlib.Path(args.results_dir), args.walk_subdirs, args.name_suffix)
+    raw_results = pathlib.Path(args.results)
+    if not args.key_sizes:
+        print("Parsing benchmark results... ")
+        results = recursive_parse(raw_results, args.walk_subdirs, args.name_suffix)
+    else:
+        print("Parsing key sizes results... ")
+        results = parse_key_sizes(raw_results)
     print("Parsing results done")
 
     output_file = pathlib.Path(args.output_file)

tfhe/Cargo.toml

@@ -129,6 +129,14 @@ required-features = ["shortint", "internal-keycache"]
 name = "generates_test_keys"
 required-features = ["shortint", "internal-keycache"]
 
+[[example]]
+name = "boolean_key_sizes"
+required-features = ["boolean", "internal-keycache"]
+
+[[example]]
+name = "shortint_key_sizes"
+required-features = ["shortint", "internal-keycache"]
+
 [[example]]
 name = "micro_bench_and"
 required-features = ["boolean"]

tfhe/examples/boolean_key_sizes.rs

@@ -0,0 +1,63 @@
+use std::fs::{File, OpenOptions};
+use std::io::Write;
+use std::path::Path;
+use tfhe::boolean::parameters::{DEFAULT_PARAMETERS, TFHE_LIB_PARAMETERS};
+use tfhe::boolean::{client_key, server_key};
+
+fn write_result(file: &mut File, name: &str, value: usize) {
+    let line = format!("{name},{value}\n");
+    let error_message = format!("cannot write {name} result into file");
+    file.write_all(line.as_bytes()).expect(&error_message);
+}
+
+fn client_server_key_sizes(results_file: &Path) {
+    let boolean_params_vec = vec![
+        (DEFAULT_PARAMETERS, "default"),
+        (TFHE_LIB_PARAMETERS, "tfhe_lib"),
+    ];
+    File::create(results_file).expect("create results file failed");
+    let mut file = OpenOptions::new()
+        .append(true)
+        .open(results_file)
+        .expect("cannot open results file");
+
+    println!("Generating boolean (ClientKey, ServerKey)");
+    for (i, (params, name)) in boolean_params_vec.iter().enumerate() {
+        println!(
+            "Generating [{} / {}] : {}",
+            i + 1,
+            boolean_params_vec.len(),
+            name
+        );
+
+        let cks = client_key::ClientKey::new(params);
+        let sks = server_key::ServerKey::new(&cks);
+        let ksk_size = sks.key_switching_key_size_bytes();
+        write_result(&mut file, &format!("boolean_{}_{}", name, "ksk"), ksk_size);
+        println!(
+            "Element in KSK: {}, size in bytes: {}",
+            sks.key_switching_key_size_elements(),
+            ksk_size,
+        );
+
+        let bsk_size = sks.bootstrapping_key_size_bytes();
+        write_result(&mut file, &format!("boolean_{}_{}", name, "bsk"), bsk_size);
+        println!(
+            "Element in BSK: {}, size in bytes: {}",
+            sks.bootstrapping_key_size_elements(),
+            bsk_size,
+        );
+    }
+}
+
+fn main() {
+    let work_dir = std::env::current_dir().unwrap();
+    println!("work_dir: {}", std::env::current_dir().unwrap().display());
+    // Change workdir so that the location of the keycache matches the one for tests
+    let mut new_work_dir = work_dir;
+    new_work_dir.push("tfhe");
+    std::env::set_current_dir(new_work_dir).unwrap();
+
+    let results_file = Path::new("boolean_key_sizes.csv");
+    client_server_key_sizes(results_file)
+}

tfhe/examples/shortint_key_sizes.rs

@@ -0,0 +1,82 @@
+use std::fs::{File, OpenOptions};
+use std::io::Write;
+use std::path::Path;
+use tfhe::shortint::keycache::{NamedParam, KEY_CACHE};
+use tfhe::shortint::parameters::{
+    PARAM_MESSAGE_1_CARRY_1, PARAM_MESSAGE_2_CARRY_2, PARAM_MESSAGE_3_CARRY_3,
+    PARAM_MESSAGE_4_CARRY_4,
+};
+
+fn write_result(file: &mut File, name: &str, value: usize) {
+    let line = format!("{name},{value}\n");
+    let error_message = format!("cannot write {name} result into file");
+    file.write_all(line.as_bytes()).expect(&error_message);
+}
+
+fn client_server_key_sizes(results_file: &Path) {
+    let shortint_params_vec = vec![
+        PARAM_MESSAGE_1_CARRY_1,
+        PARAM_MESSAGE_2_CARRY_2,
+        PARAM_MESSAGE_3_CARRY_3,
+        PARAM_MESSAGE_4_CARRY_4,
+    ];
+    File::create(results_file).expect("create results file failed");
+    let mut file = OpenOptions::new()
+        .append(true)
+        .open(results_file)
+        .expect("cannot open results file");
+
+    println!("Generating shortint (ClientKey, ServerKey)");
+    for (i, params) in shortint_params_vec.iter().enumerate() {
+        println!(
+            "Generating [{} / {}] : {}",
+            i + 1,
+            shortint_params_vec.len(),
+            params.name()
+        );
+
+        let keys = KEY_CACHE.get_from_param(*params);
+
+        // Client keys don't have public access to members, but the keys in there are small anyways
+        // let cks = keys.client_key();
+        let sks = keys.server_key();
+        let ksk_size = sks.key_switching_key_size_bytes();
+        write_result(
+            &mut file,
+            &format!("shortint_{}_ksk", params.name().to_lowercase()),
+            ksk_size,
+        );
+        println!(
+            "Element in KSK: {}, size in bytes: {}",
+            sks.key_switching_key_size_elements(),
+            ksk_size,
+        );
+
+        let bsk_size = sks.bootstrapping_key_size_bytes();
+        write_result(
+            &mut file,
+            &format!("shortint_{}_bsk", params.name().to_lowercase()),
+            bsk_size,
+        );
+        println!(
+            "Element in BSK: {}, size in bytes: {}",
+            sks.bootstrapping_key_size_elements(),
+            bsk_size,
+        );
+
+        // Clear keys as we go to avoid filling the RAM
+        KEY_CACHE.clear_in_memory_cache()
+    }
+}
+
+fn main() {
+    let work_dir = std::env::current_dir().unwrap();
+    println!("work_dir: {}", std::env::current_dir().unwrap().display());
+    // Change workdir so that the location of the keycache matches the one for tests
+    let mut new_work_dir = work_dir;
+    new_work_dir.push("tfhe");
+    std::env::set_current_dir(new_work_dir).unwrap();
+
+    let results_file = Path::new("shortint_key_sizes.csv");
+    client_server_key_sizes(results_file)
+}

tfhe/src/boolean/engine/bootstrapping.rs

@@ -84,6 +84,24 @@ pub struct ServerKey {
     pub(crate) key_switching_key: LweKeyswitchKeyOwned<u32>,
 }
 
+impl ServerKey {
+    pub fn bootstrapping_key_size_elements(&self) -> usize {
+        self.bootstrapping_key.as_view().data().as_ref().len()
+    }
+
+    pub fn bootstrapping_key_size_bytes(&self) -> usize {
+        self.bootstrapping_key_size_elements() * std::mem::size_of::<concrete_fft::c64>()
+    }
+
+    pub fn key_switching_key_size_elements(&self) -> usize {
+        self.key_switching_key.as_ref().len()
+    }
+
+    pub fn key_switching_key_size_bytes(&self) -> usize {
+        self.key_switching_key_size_elements() * std::mem::size_of::<u64>()
+    }
+}
+
 /// Perform ciphertext bootstraps on the CPU
 pub(crate) struct Bootstrapper {
     memory: Memory,

tfhe/src/shortint/server_key/mod.rs

@@ -534,6 +534,22 @@ impl ServerKey {
             engine.create_trivial_assign(self, ct, value).unwrap()
         })
     }
+
+    pub fn bootstrapping_key_size_elements(&self) -> usize {
+        self.bootstrapping_key.as_view().data().as_ref().len()
+    }
+
+    pub fn bootstrapping_key_size_bytes(&self) -> usize {
+        self.bootstrapping_key_size_elements() * std::mem::size_of::<concrete_fft::c64>()
+    }
+
+    pub fn key_switching_key_size_elements(&self) -> usize {
+        self.key_switching_key.as_ref().len()
+    }
+
+    pub fn key_switching_key_size_bytes(&self) -> usize {
+        self.key_switching_key_size_elements() * std::mem::size_of::<u64>()
+    }
 }
 
 #[derive(Serialize, Deserialize)]