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,
    CompressedFheUint8,
    FheUint8,
    FheUint32,
    CompactFheUint32,
    CompactFheUint32List,
    CompressedFheUint128,
    FheUint128,
    CompressedFheUint256,
    FheUint256,
    CompactFheUint256,
    CompactFheUint256List,
} from "./pkg/tfhe.js";

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()
        .use_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)
}

const U32_MAX = 4294967295;

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

    console.time('CompactPublicKey Gen')
    let publicKey = TfheCompactPublicKey.new(clientKey);
    console.timeEnd('CompactPublicKey Gen')

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

    let values = [0, 1, 2394, U32_MAX];

    console.time('CompactFheUint32List Encrypt')
    let compact_list = CompactFheUint32List.encrypt_with_compact_public_key(values, publicKey);
    console.timeEnd('CompactFheUint32List Encrypt')

    {
        console.time('CompactFheUint32List Expand')
        let encrypted_list = compact_list.expand();
        console.timeEnd('CompactFheUint32List Expand')

        assert_eq(encrypted_list.length, values.length);

        for (let i = 0; i < values.length; i++) {
            let decrypted = encrypted_list[i].decrypt(clientKey);
            assert_eq(decrypted, values[i]);
        }
    }

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

    let deserialized_list = CompactFheUint32List.deserialize(serialized_list);
    let encrypted_list = deserialized_list.expand();
    assert_eq(encrypted_list.length, values.length);

    for (let i = 0; i < values.length; i++) {
        let decrypted = encrypted_list[i].decrypt(clientKey);
        assert_eq(decrypted, values[i]);
    }
}

async function compactPublicKeyTest32BitBig() {
    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 compactPublicKeyTest32BitOnConfig(config)
}

async function compactPublicKeyTest32BitSmall() {
    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 compactPublicKeyTest32BitOnConfig(config)
}

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];

    // Encrypt compact CT list for serialization for later
    console.time('CompactFheUint32List Encrypt')
    let compact_list = CompactFheUint32List.encrypt_with_compact_public_key(values, publicKey);
    console.timeEnd('CompactFheUint32List Encrypt')

    // Bench the encryption for bench_loops iterations
    start = performance.now();
    for (let i = 0; i < bench_loops; i++) {
        let _ = CompactFheUint32List.encrypt_with_compact_public_key(values, publicKey);
    }
    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"
    );
}

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

    console.time('CompactPublicKey Gen')
    let publicKey = TfheCompactPublicKey.new(clientKey);
    console.timeEnd('CompactPublicKey Gen')

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

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

    console.time('CompactFheUint256List Encrypt')
    let compact_list = CompactFheUint256List.encrypt_with_compact_public_key(values, publicKey);
    console.timeEnd('CompactFheUint256List Encrypt')

    {
        console.time('CompactFheUint256List Expand')
        let encrypted_list = compact_list.expand();
        console.timeEnd('CompactFheUint256List Expand')

        assert_eq(encrypted_list.length, values.length);

        for (let i = 0; i < values.length; i++) {
            let decrypted = encrypted_list[i].decrypt(clientKey);
            assert_eq(decrypted, values[i]);
        }
    }

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

    let deserialized_list = CompactFheUint256List.deserialize(serialized_list);
    let encrypted_list = deserialized_list.expand();
    assert_eq(encrypted_list.length, values.length);

    for (let i = 0; i < values.length; i++) {
        let decrypted = encrypted_list[i].decrypt(clientKey);
        assert_eq(decrypted, values[i]);
    }
}

async function compactPublicKeyTest256BitBig() {
    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 compactPublicKeyTest256BitOnConfig(config)
}

async function compactPublicKeyTest256BitSmall() {
    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 compactPublicKeyTest256BitOnConfig(config)
}


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 values = [0, 1, 2394, U32_MAX].map((e) => BigInt(e));

    console.time('CompactFheUint256List Encrypt')
    let compact_list = CompactFheUint256List.encrypt_with_compact_public_key(values, publicKey);
    console.timeEnd('CompactFheUint256List Encrypt')

    {
        console.time('CompactFheUint256List Expand')
        let encrypted_list = compact_list.expand();
        console.timeEnd('CompactFheUint256List Expand')

        assert_eq(encrypted_list.length, values.length);

        for (let i = 0; i < values.length; i++) {
            let decrypted = encrypted_list[i].decrypt(clientKey);
            assert_eq(decrypted, values[i]);
        }
    }

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

    let deserialized_list = CompactFheUint256List.deserialize(serialized_list);
    let encrypted_list = deserialized_list.expand();
    assert_eq(encrypted_list.length, values.length);

    for (let i = 0; i < values.length; i++) {
        let decrypted = encrypted_list[i].decrypt(clientKey);
        assert_eq(decrypted, values[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));

    // Encrypt compact CT list for serialization for later
    console.time('CompactFheUint256List Encrypt')
    let compact_list = CompactFheUint256List.encrypt_with_compact_public_key(values, publicKey);
    console.timeEnd('CompactFheUint256List Encrypt')

    // Bench the encryption for bench_loops iterations
    start = performance.now();
    for (let i = 0; i < bench_loops; i++) {
        let _ = CompactFheUint256List.encrypt_with_compact_public_key(values, publicKey);
    }
    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 main() {
    await init()
    await initThreadPool(navigator.hardwareConcurrency);
    await init_panic_hook();

    return Comlink.proxy({
        publicKeyTest,
        compressedPublicKeyTest,
        compactPublicKeyTest32BitSmall,
        compactPublicKeyTest32BitBig,
        compactPublicKeyTest256BitSmall,
        compactPublicKeyTest256BitBig,
        compressedCompactPublicKeyTest256BitSmall,
        compressedCompactPublicKeyTest256BitBig,
        compactPublicKeyBench32BitBig,
        compactPublicKeyBench32BitSmall,
        compactPublicKeyBench256BitBig,
        compactPublicKeyBench256BitSmall,
        compressedServerKeyBenchMessage1Carry1,
        compressedServerKeyBenchMessage2Carry2,
    })
}

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