tfhe-rs/tfhe/web_wasm_parallel_tests/worker.js

import * as Comlink from "comlink";
import init, {
  initThreadPool,
  init_panic_hook,
  ShortintParametersName,
  ShortintParameters,
  TfheClientKey,
  TfhePublicKey,
  TfheCompressedPublicKey,
  TfheCompressedServerKey,
  TfheCompressedCompactPublicKey,
  TfheCompactPublicKey,
  TfheConfigBuilder,
  FheUint8,
  ZkComputeLoad,
  CompactPkeCrs,
  CompactCiphertextList,
  ProvenCompactCiphertextList,
  Shortint,
  ShortintEncryptionKeyChoice,
} from "./pkg/tfhe.js";

const U32_MAX = 4294967295;

function assert(cond, text) {
  if (cond) return;
  if (console.assert.useDebugger) debugger;
  throw new Error(text || "Assertion failed!");
}

function assert_eq(a, b, text) {
  if (a === b) return;
  if (console.assert.useDebugger) debugger;
  throw new Error(text || `Equality assertion failed!: ${a} != ${b}`);
}

function append_param_name(bench_results, params_name) {
  let results = {};
  for (const bench_name in bench_results) {
    results[`${bench_name}_${params_name}`] = bench_results[bench_name];
  }
  return results;
}

async function compressedPublicKeyTest() {
  let config = TfheConfigBuilder.default_with_small_encryption().build();

  console.time("ClientKey Gen");
  let clientKey = TfheClientKey.generate(config);
  console.timeEnd("ClientKey Gen");

  console.time("CompressedPublicKey Gen");
  let compressedPublicKey = TfheCompressedPublicKey.new(clientKey);
  console.timeEnd("CompressedPublicKey Gen");

  let data = compressedPublicKey.serialize();
  console.log("CompressedPublicKey size:", data.length);

  console.time("CompressedPublicKey Decompression");
  let publicKey = compressedPublicKey.decompress();
  console.timeEnd("CompressedPublicKey Decompression");

  console.time("FheUint8 encrypt with CompressedPublicKey");
  let encrypted = FheUint8.encrypt_with_public_key(255, publicKey);
  console.timeEnd("FheUint8 encrypt with CompressedPublicKey");

  let ser = encrypted.serialize();
  console.log("Ciphertext Size", ser.length);

  let decrypted = encrypted.decrypt(clientKey);
  assert_eq(decrypted, 255);
}

async function publicKeyTest() {
  let config = TfheConfigBuilder.default_with_small_encryption().build();

  console.time("ClientKey Gen");
  let clientKey = TfheClientKey.generate(config);
  console.timeEnd("ClientKey Gen");

  console.time("PublicKey Gen");
  let publicKey = TfhePublicKey.new(clientKey);
  console.timeEnd("PublicKey Gen");

  console.time("FheUint8 encrypt with PublicKey");
  let encrypted = FheUint8.encrypt_with_public_key(255, publicKey);
  console.timeEnd("FheUint8 encrypt with PublicKey");

  let ser = encrypted.serialize();
  console.log("Ciphertext Size", ser.length);

  let decrypted = encrypted.decrypt(clientKey);
  assert_eq(decrypted, 255);
}

async function compactPublicKeyBench32BitOnConfig(config) {
  const bench_loops = 100;
  let bench_results = {};

  console.time("ClientKey Gen");
  let clientKey = TfheClientKey.generate(config);
  console.timeEnd("ClientKey Gen");

  // Generate PK for encryption for later
  console.time("CompactPublicKey Gen");
  let publicKey = TfheCompactPublicKey.new(clientKey);
  console.timeEnd("CompactPublicKey Gen");

  // Bench the pk generation for bench_loops iterations
  let start = performance.now();
  for (let i = 0; i < bench_loops; i++) {
    let _ = TfheCompactPublicKey.new(clientKey);
  }
  let end = performance.now();
  const timing_1 = (end - start) / bench_loops;
  console.log("CompactPublicKey Gen bench: ", timing_1, " ms");
  bench_results["compact_public_key_gen_32bit_mean"] = timing_1;

  let values = [0, 1, 2, 2394, U32_MAX].map(BigInt);

  // Bench the encryption for bench_loops iterations
  start = performance.now();
  let compact_list;
  for (let i = 0; i < bench_loops; i++) {
    let builder = CompactCiphertextList.builder(publicKey);
    for (let value of values) {
      builder.push_u256(value);
    }
    compact_list = builder.build();
  }
  end = performance.now();
  const timing_2 = (end - start) / bench_loops;
  console.log("CompactFheUint32List Encrypt bench: ", timing_2, " ms");
  bench_results["compact_fheunit32_list_encrypt_mean"] = timing_2;

  let serialized_list = compact_list.serialize();
  console.log("Serialized CompactFheUint32List size: ", serialized_list.length);

  // Bench the serialization for bench_loops iterations
  start = performance.now();
  for (let i = 0; i < bench_loops; i++) {
    let _ = compact_list.serialize();
  }
  end = performance.now();
  const timing_3 = (end - start) / bench_loops;
  console.log("CompactFheUint32List serialization bench: ", timing_3, " ms");
  bench_results["compact_fheunit32_list_serialization_mean"] = timing_3;

  return bench_results;
}

async function compactPublicKeyBench32BitBig() {
  const block_params = new ShortintParameters(
    ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_KS_PBS,
  );
  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();
  return append_param_name(
    await compactPublicKeyBench32BitOnConfig(config),
    "PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_KS_PBS",
  );
}

async function compactPublicKeyBench32BitSmall() {
  const block_params = new ShortintParameters(
    ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_PBS_KS,
  );
  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();
  return append_param_name(
    await compactPublicKeyBench32BitOnConfig(config),
    "PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_PBS_KS",
  );
}

function generateRandomBigInt(bitLen) {
  let result = BigInt(0);
  for (let i = 0; i < bitLen; i++) {
    result << 1n;
    result |= BigInt(Math.random() < 0.5);
  }
  return result;
}

async function compressedCompactPublicKeyTest256BitOnConfig(config) {
  console.time("ClientKey Gen");
  let clientKey = TfheClientKey.generate(config);
  console.timeEnd("ClientKey Gen");

  console.time("CompressedCompactPublicKey Gen");
  let publicKey = TfheCompressedCompactPublicKey.new(clientKey);
  console.timeEnd("CompressedCompactPublicKey Gen");

  let serialized_pk = publicKey.serialize();
  console.log(
    "Serialized CompressedCompactPublicKey size: ",
    serialized_pk.length,
  );

  console.time("CompressedCompactPublicKey Decompression");
  publicKey = publicKey.decompress();
  console.timeEnd("CompressedCompactPublicKey Decompression");

  let clear_u2 = 3;
  let clear_i32 = -3284;
  let clear_bool = true;
  let clear_u256 = generateRandomBigInt(256);

  let builder = CompactCiphertextList.builder(publicKey);
  builder.push_u2(clear_u2);
  builder.push_i32(clear_i32);
  builder.push_boolean(clear_bool);
  builder.push_u256(clear_u256);

  let num_bits_encrypted = 2 + 4 + 1 + 256;
  console.log("Numb bits in compact list: ", num_bits_encrypted);

  console.time("CompactCiphertextList Encrypt");
  let list = builder.build();
  console.timeEnd("CompactCiphertextList Encrypt");

  let serialized = list.safe_serialize(BigInt(10000000));
  console.log("Serialized CompactCiphertextList size: ", serialized.length);
  let deserialized = CompactCiphertextList.safe_deserialize(
    serialized,
    BigInt(10000000),
  );

  let expander = deserialized.expand();

  assert_eq(expander.get_uint2(0).decrypt(clientKey), clear_u2);

  assert_eq(expander.get_int32(1).decrypt(clientKey), clear_i32);

  assert_eq(expander.get_bool(2).decrypt(clientKey), clear_bool);

  assert_eq(expander.get_uint256(3).decrypt(clientKey), clear_u256);
}

async function compactPublicKeyZeroKnowledge() {
  let block_params = Shortint.new_parameters(
    888,
    2,
    2048,
    Shortint.try_new_t_uniform(45),
    Shortint.try_new_t_uniform(3),
    23,
    1,
    4,
    4,
    4,
    4,
    5,
    -64.105,
    64,
    ShortintEncryptionKeyChoice.Big,
  );

  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();

  let clientKey = TfheClientKey.generate(config);
  let publicKey = TfheCompactPublicKey.new(clientKey);

  console.log("Start CRS generation");
  console.time("CRS generation");
  let crs = CompactPkeCrs.from_config(config, 4 * 64);
  console.timeEnd("CRS generation");
  let public_params = crs.public_params();

  {
    let input = generateRandomBigInt(64);
    let start = performance.now();

    let builder = CompactCiphertextList.builder(publicKey);
    builder.push_u64(input);
    let list = builder.build_with_proof(public_params, ZkComputeLoad.Proof);
    let end = performance.now();
    console.log(
      "Time to encrypt + prove CompactFheUint64: ",
      end - start,
      " ms",
    );

    let bytes = list.serialize();
    console.log("CompactCiphertextList size:", bytes.length);

    start = performance.now();
    let expander = list.verify_and_expand(public_params, publicKey);
    end = performance.now();
    console.log(
      "Time to verify + expand CompactFheUint64: ",
      end - start,
      " ms",
    );

    assert_eq(expander.get_uint64(0).decrypt(clientKey), input);
  }

  {
    let inputs = [
      generateRandomBigInt(64),
      generateRandomBigInt(64),
      generateRandomBigInt(64),
      generateRandomBigInt(64),
    ];
    let start = performance.now();
    let builder = CompactCiphertextList.builder(publicKey);
    for (let input of inputs) {
      builder.push_u64(input);
    }
    let encrypted = builder.build_with_proof(
      public_params,
      ZkComputeLoad.Proof,
    );
    let end = performance.now();
    console.log(
      "Time to encrypt + prove CompactFheUint64List of 4: ",
      end - start,
      " ms",
    );

    start = performance.now();
    let expander = encrypted.verify_and_expand(public_params, publicKey);
    end = performance.now();
    console.log(
      "Time to verify + expand CompactFheUint64: ",
      end - start,
      " ms",
    );

    for (let i = 0; i < inputs.length; i++) {
      assert_eq(expander.get_uint64(i).decrypt(clientKey), inputs[i]);
    }
  }
}

async function compressedCompactPublicKeyTest256BitBig() {
  const block_params = new ShortintParameters(
    ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_KS_PBS,
  );
  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();
  await compressedCompactPublicKeyTest256BitOnConfig(config);
}

async function compressedCompactPublicKeyTest256BitSmall() {
  const block_params = new ShortintParameters(
    ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_PBS_KS,
  );
  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();
  await compressedCompactPublicKeyTest256BitOnConfig(config);
}

async function compactPublicKeyBench256BitOnConfig(config) {
  const bench_loops = 100;
  let bench_results = {};

  console.time("ClientKey Gen");
  let clientKey = TfheClientKey.generate(config);
  console.timeEnd("ClientKey Gen");

  // Generate PK for encryption for later
  console.time("CompactPublicKey Gen");
  let publicKey = TfheCompactPublicKey.new(clientKey);
  console.timeEnd("CompactPublicKey Gen");

  // Bench the pk generation for bench_loops iterations
  let start = performance.now();
  for (let i = 0; i < bench_loops; i++) {
    let _ = TfheCompactPublicKey.new(clientKey);
  }
  let end = performance.now();
  const timing_1 = (end - start) / bench_loops;
  console.log("CompactPublicKey Gen bench: ", timing_1, " ms");
  bench_results["compact_public_key_gen_256bit_mean"] = timing_1;

  let values = [0, 1, 2, 2394, U32_MAX].map((e) => BigInt(e));

  // Bench the encryption for bench_loops iterations
  start = performance.now();
  let compact_list;
  for (let i = 0; i < bench_loops; i++) {
    console.time("CompactFheUint256List Encrypt");
    let builder = CompactCiphertextList.builder(publicKey);
    for (let value of values) {
      builder.push_u256(value);
    }
    compact_list = builder.build();
    console.timeEnd("CompactFheUint256List Encrypt");
  }
  end = performance.now();
  const timing_2 = (end - start) / bench_loops;
  console.log("CompactFheUint256List Encrypt bench: ", timing_2, " ms");
  bench_results["compact_fheunit256_list_encrypt_mean"] = timing_2;

  let serialized_list = compact_list.serialize();
  console.log(
    "Serialized CompactFheUint256List size: ",
    serialized_list.length,
  );

  // Bench the serialization for bench_loops iterations
  start = performance.now();
  for (let i = 0; i < bench_loops; i++) {
    let _ = compact_list.serialize();
  }
  end = performance.now();
  const timing_3 = (end - start) / bench_loops;
  console.log("CompactFheUint256List serialization bench: ", timing_3, " ms");
  bench_results["compact_fheunit256_list_serialization_mean"] = timing_3;

  return bench_results;
}

async function compactPublicKeyBench256BitBig() {
  const block_params = new ShortintParameters(
    ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_KS_PBS,
  );
  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();
  return append_param_name(
    await compactPublicKeyBench256BitOnConfig(config),
    "PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_KS_PBS",
  );
}

async function compactPublicKeyBench256BitSmall() {
  const block_params = new ShortintParameters(
    ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_PBS_KS,
  );
  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();
  return append_param_name(
    await compactPublicKeyBench256BitOnConfig(config),
    "PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_PBS_KS",
  );
}

async function compressedServerKeyBenchConfig(config) {
  const bench_loops = 35;
  let bench_results = {};

  console.log("Begin benchmarks"); // DEBUG
  let clientKey = TfheClientKey.generate(config);

  // Bench the sk generation for bench_loops iterations
  let start = performance.now();
  for (let i = 0; i < bench_loops; i++) {
    let _ = TfheCompressedServerKey.new(clientKey);
  }
  let end = performance.now();
  const timing_1 = (end - start) / bench_loops;
  console.log("CompressedServerKey Gen bench: ", timing_1, " ms");
  bench_results["compressed_server_key_gen_mean"] = timing_1;

  let serverKey = TfheCompressedServerKey.new(clientKey);
  let serialized_key = serverKey.serialize();
  console.log("Serialized ServerKey size: ", serialized_key.length);

  // Bench the serialization for bench_loops iterations
  start = performance.now();
  for (let i = 0; i < bench_loops; i++) {
    let _ = serverKey.serialize();
  }
  end = performance.now();
  const timing_2 = (end - start) / bench_loops;
  console.log("CompressedServerKey serialization bench: ", timing_2, " ms");
  bench_results["compressed_server_key_serialization_mean"] = timing_2;

  return bench_results;
}

async function compressedServerKeyBenchMessage1Carry1() {
  const block_params = new ShortintParameters(
    ShortintParametersName.PARAM_MESSAGE_1_CARRY_1_KS_PBS,
  );
  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();
  return append_param_name(
    await compressedServerKeyBenchConfig(config),
    "PARAM_MESSAGE_1_CARRY_1_KS_PBS",
  );
}

async function compressedServerKeyBenchMessage2Carry2() {
  const block_params = new ShortintParameters(
    ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_KS_PBS,
  );
  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();
  return append_param_name(
    await compressedServerKeyBenchConfig(config),
    "PARAM_MESSAGE_2_CARRY_2_KS_PBS",
  );
}

async function compactPublicKeyZeroKnowledgeBench() {
  // This parameters set reproduce PARAM_MESSAGE_2_CARRY_2_COMPACT_PK_KS_PBS_TUNIFORM_2M64
  let block_params = Shortint.new_parameters(
    888,
    2,
    2048,
    Shortint.try_new_t_uniform(45),
    Shortint.try_new_t_uniform(3),
    23,
    1,
    4,
    4,
    4,
    4,
    5,
    -64.105,
    64,
    ShortintEncryptionKeyChoice.Big,
  );

  let config = TfheConfigBuilder.default()
    .use_custom_parameters(block_params)
    .build();

  let clientKey = TfheClientKey.generate(config);
  let publicKey = TfheCompactPublicKey.new(clientKey);

  console.log("Start CRS generation");
  console.time("CRS generation");
  let crs = CompactPkeCrs.from_config(config, 4 * 64);
  console.timeEnd("CRS generation");
  let public_params = crs.public_params();

  const bench_loops = 4; // The computation is expensive
  let bench_results = {};
  let load_choices = [ZkComputeLoad.Proof, ZkComputeLoad.Verify];
  const load_to_str = {
    [ZkComputeLoad.Proof]: "compute_load_proof",
    [ZkComputeLoad.Verify]: "compute_load_verify",
  };
  for (const loadChoice of load_choices) {
    let timing = 0;
    for (let i = 0; i < bench_loops; i++) {
      let input = generateRandomBigInt(64);

      const start = performance.now();
      let builder = ProvenCompactCiphertextList.builder(publicKey);
      builder.push_u64(input);
      let list = builder.build_with_proof(public_params);

      const end = performance.now();
      timing += end - start;
    }
    const mean = timing / bench_loops;

    const bench_str =
      "compact_fhe_uint64_proven_encryption_" +
      load_to_str[loadChoice] +
      "_mean";
    console.log(bench_str, ": ", mean, " ms");
    bench_results["compact_fhe_uint64_proven_encryption_"] = mean;
  }

  return bench_results;
}

async function main() {
  await init();
  await initThreadPool(navigator.hardwareConcurrency);
  await init_panic_hook();

  return Comlink.proxy({
    publicKeyTest,
    compressedPublicKeyTest,
    compressedCompactPublicKeyTest256BitSmall,
    compressedCompactPublicKeyTest256BitBig,
    compactPublicKeyZeroKnowledge,
    compactPublicKeyBench32BitBig,
    compactPublicKeyBench32BitSmall,
    compactPublicKeyBench256BitBig,
    compactPublicKeyBench256BitSmall,
    compressedServerKeyBenchMessage1Carry1,
    compressedServerKeyBenchMessage2Carry2,
    compactPublicKeyZeroKnowledgeBench,
  });
}

Comlink.expose({
  demos: main(),
});