From ddf78278722b1b3da4117ed33f13a1e968a87246 Mon Sep 17 00:00:00 2001 From: youben11 Date: Thu, 14 Nov 2024 14:10:40 +0100 Subject: [PATCH] draft ML example --- .../examples/tfhers-ml/README.md | 91 ++++++++ .../examples/tfhers-ml/example.py | 197 ++++++++++++++++++ .../examples/tfhers-ml/test.sh | 19 ++ .../examples/tfhers-ml/tfhers_params.json | 1 + 4 files changed, 308 insertions(+) create mode 100644 frontends/concrete-python/examples/tfhers-ml/README.md create mode 100644 frontends/concrete-python/examples/tfhers-ml/example.py create mode 100755 frontends/concrete-python/examples/tfhers-ml/test.sh create mode 100644 frontends/concrete-python/examples/tfhers-ml/tfhers_params.json diff --git a/frontends/concrete-python/examples/tfhers-ml/README.md b/frontends/concrete-python/examples/tfhers-ml/README.md new file mode 100644 index 0000000000..9d4971dce9 --- /dev/null +++ b/frontends/concrete-python/examples/tfhers-ml/README.md @@ -0,0 +1,91 @@ +# TFHE-rs Interoperability Example + +This is the full execution for the example explained in the [TFHE-rs Interoperability Guide](../../../../docs/guides/tfhers) (use case 1). You can find the TFHE-rs code [here](../../tests/tfhers-utils/src/main.rs), while the Python code is under this direcotry [here](example.py). Both are CLI tools, so that we can execute the example step by step. You can refer to the code at every step to see how it's implemented. + +## Make tmpdir + +We want to setup a temporary working directory first + +```sh +export TDIR=`mktemp -d` +``` + +## KeyGen + +First we need to build the TFHE-rs utility in [this directory](../../tests/tfhers-utils/) by running + +```sh +cd ../../tests/tfhers-utils/ +make build +cd - +``` + +Then we can generate keys in two different ways. You only need to run one of the following methods + +#### Generate the Secret Key in Concrete + +We start by doing keygen in Concrete + +```sh +python example.py keygen -o $TDIR/concrete_sk -k $TDIR/concrete_keyset +``` + +Then we do a partial keygen in TFHE-rs + +```sh +../../tests/tfhers-utils/target/release/tfhers_utils keygen --lwe-sk $TDIR/concrete_sk --output-lwe-sk $TDIR/tfhers_sk -c $TDIR/tfhers_client_key -s $TDIR/tfhers_server_key +``` + +#### Generate the Secret Key in TFHE-rs + +We start by doing keygen in TFHE-rs + +```sh +../../tests/tfhers-utils/target/release/tfhers_utils keygen --output-lwe-sk $TDIR/tfhers_sk -c $TDIR/tfhers_client_key -s $TDIR/tfhers_server_key +``` + +Then we do a partial keygen in Concrete. + +```sh +python example.py keygen -s $TDIR/tfhers_sk -o $TDIR/concrete_sk -k $TDIR/concrete_keyset +``` + +## Encrypt in TFHE-rs + +```sh +../../tests/tfhers-utils/target/release/tfhers_utils encrypt-with-key --value 162 --ciphertext $TDIR/tfhers_ct_1 --client-key $TDIR/tfhers_client_key +../../tests/tfhers-utils/target/release/tfhers_utils encrypt-with-key --value 73 --ciphertext $TDIR/tfhers_ct_2 --client-key $TDIR/tfhers_client_key +``` + +{% hint style="info" %} + +If you have tensor inputs, then you can encrypt by passing your flat tensor in `--value`. Concrete will take care of reshaping the values to the corresponding shape. For example `--value=1,2,3,4` can represent a 2 by 2 tensor, or a flat vector of 4 values. + +{% endhint %} + +## Compute in TFHE-rs + +```sh +# encrypt value to add first +../../tests/tfhers-utils/target/release/tfhers_utils encrypt-with-key --value 9 --ciphertext $TDIR/tfhers_ct_inc --client-key $TDIR/tfhers_client_key +# add two ciphertexts +../../tests/tfhers-utils/target/release/tfhers_utils add --server-key $TDIR/tfhers_server_key --cts $TDIR/tfhers_ct_2 $TDIR/tfhers_ct_inc --output-ct $TDIR/tfhers_ct_2 +``` + +## Run in Concrete + +```sh +python example.py run -k $TDIR/concrete_keyset -c1 $TDIR/tfhers_ct_1 -c2 $TDIR/tfhers_ct_2 -o $TDIR/tfhers_ct_out +``` + +## Decrypt in TFHE-rs + +```sh +../../tests/tfhers-utils/target/release/tfhers_utils decrypt-with-key --ciphertext $TDIR/tfhers_ct_out --client-key $TDIR/tfhers_client_key +``` + +## Clean tmpdir + +```sh +rm -rf $TDIR +``` diff --git a/frontends/concrete-python/examples/tfhers-ml/example.py b/frontends/concrete-python/examples/tfhers-ml/example.py new file mode 100644 index 0000000000..54ae259620 --- /dev/null +++ b/frontends/concrete-python/examples/tfhers-ml/example.py @@ -0,0 +1,197 @@ +import os +from functools import partial + +import click +import numpy as np + +from concrete import fhe +from concrete.fhe import tfhers + +### Options ########################### +# These parameters were saved by running the tfhers_utils utility: +# tfhers_utils save-params tfhers_params.json +TFHERS_PARAMS_FILE = "tfhers_params.json" +FHEUINT_PRECISION = 8 +IS_SIGNED = True +####################################### + +tfhers_type = tfhers.get_type_from_params( + TFHERS_PARAMS_FILE, + is_signed=IS_SIGNED, + precision=FHEUINT_PRECISION, +) +tfhers_int = partial(tfhers.TFHERSInteger, tfhers_type) + + +q_weights = np.array( + [ + [-81], + [-95], + [-51], + [-77], + [-64], + [-64], + [-128], + [127], + [-122], + [-81], + [-96], + [-93], + [-63], + [-50], + [-104], + [-99], + [-112], + [-46], + [-106], + [-42], + [-56], + [-46], + [-67], + [-116], + [-107], + [-75], + [-105], + [-109], + [-88], + [-80], + ] +) +q_weights = np.ones_like(q_weights) +q_bias = np.array([[680]]) +weight_quantizer_zero_point = -74 + + +def ml_inference(q_X: np.ndarray) -> np.ndarray: + # Quantizing weights and inputs makes an additional term appear in the inference function + y_pred = q_X @ q_weights - weight_quantizer_zero_point * np.sum(q_X, axis=1, keepdims=True) + y_pred += q_bias + return y_pred + + +def compute(tfhers_x): + ####### TFHE-rs to Concrete ######### + + # x and y are supposed to be TFHE-rs values. + # to_native will use type information from x and y to do + # a correct conversion from TFHE-rs to Concrete + concrete_x = tfhers.to_native(tfhers_x) + ####### TFHE-rs to Concrete ######### + + ####### Concrete Computation ######## + concrete_res = ml_inference(concrete_x) + ####### Concrete Computation ######## + + ####### Concrete to TFHE-rs ######### + tfhers_res = tfhers.from_native( + concrete_res, tfhers_type + ) # we have to specify the type we want to convert to + ####### Concrete to TFHE-rs ######### + return tfhers_res + + +def ccompilee(): + compiler = fhe.Compiler( + compute, + { + "tfhers_x": "encrypted", + }, + ) + + inputset = [ + ( + tfhers_int( + np.array( + [ + 1, + ] + * 30 + ).reshape((1, 30)) + ), + ) + ] + circuit = compiler.compile(inputset, show_graph=True, show_mlir=True) + + tfhers_bridge = tfhers.new_bridge(circuit=circuit) + return circuit, tfhers_bridge + + +@click.group() +def cli(): + pass + + +@cli.command() +@click.option("-s", "--secret-key", type=str, required=False) +@click.option("-o", "--output-secret-key", type=str, required=True) +@click.option("-k", "--concrete-keyset-path", type=str, required=True) +def keygen(output_secret_key: str, secret_key: str, concrete_keyset_path: str): + """Concrete Key Generation""" + + circuit, tfhers_bridge = ccompilee() + + if os.path.exists(concrete_keyset_path): + print(f"removing old keyset at '{concrete_keyset_path}'") + os.remove(concrete_keyset_path) + + if secret_key: + print(f"partial keygen from sk at '{secret_key}'") + # load the initial secret key to use for keygen + with open( + secret_key, + "rb", + ) as f: + buff = f.read() + input_idx_to_key = {0: buff} + tfhers_bridge.keygen_with_initial_keys(input_idx_to_key_buffer=input_idx_to_key) + else: + print("full keygen") + circuit.keygen() + + print("saving Concrete keyset") + circuit.client.keys.save(concrete_keyset_path) + print(f"saved Concrete keyset to '{concrete_keyset_path}'") + + sk: bytes = tfhers_bridge.serialize_input_secret_key(input_idx=0) + print(f"writing secret key of size {len(sk)} to '{output_secret_key}'") + with open(output_secret_key, "wb") as f: + f.write(sk) + + +@cli.command() +@click.option("-c1", "--rust-ct-1", type=str, required=True) +@click.option("-o", "--output-rust-ct", type=str, required=False) +@click.option("-k", "--concrete-keyset-path", type=str, required=True) +def run(rust_ct_1: str, output_rust_ct: str, concrete_keyset_path: str): + """Run circuit""" + circuit, tfhers_bridge = ccompilee() + + if not os.path.exists(concrete_keyset_path): + raise RuntimeError("cannot find keys, you should run keygen before") + print(f"loading keys from '{concrete_keyset_path}'") + circuit.client.keys.load(concrete_keyset_path) + + # read tfhers int from file + with open(rust_ct_1, "rb") as f: + buff = f.read() + # import fheuint8 and get its description + tfhers_uint8_x = tfhers_bridge.import_value(buff, input_idx=0) + + print("Homomorphic evaluation...") + encrypted_result = circuit.run(tfhers_uint8_x) + + if output_rust_ct: + print("exporting Rust ciphertexts") + # export fheuint8 + buff = tfhers_bridge.export_value(encrypted_result, output_idx=0) + # write it to file + with open(output_rust_ct, "wb") as f: + f.write(buff) + else: + result = circuit.decrypt(encrypted_result) + decoded = tfhers_type.decode(result) + print(f"Concrete decryption result: raw({result}), decoded({decoded})") + + +if __name__ == "__main__": + cli() diff --git a/frontends/concrete-python/examples/tfhers-ml/test.sh b/frontends/concrete-python/examples/tfhers-ml/test.sh new file mode 100755 index 0000000000..77eeccfcf5 --- /dev/null +++ b/frontends/concrete-python/examples/tfhers-ml/test.sh @@ -0,0 +1,19 @@ +#!/bin/sh + +# This file tests that the example is working + +shell_blocks=$(sed -n '/^```sh/,/^```/ p' < README.md | sed '/^```sh/d' | sed '/^```/d') + +set -e +output=$(eval "$shell_blocks" 2>&1) || echo "$output" + +result=$(echo "$output" | grep "result: " | sed 's/result: //g') + +expected="31" +if [ $result -eq $expected ] +then + exit 0 +else + echo "expected result to be $expected, but result was $result" + exit 1 +fi diff --git a/frontends/concrete-python/examples/tfhers-ml/tfhers_params.json b/frontends/concrete-python/examples/tfhers-ml/tfhers_params.json new file mode 100644 index 0000000000..bf62c959f2 --- /dev/null +++ b/frontends/concrete-python/examples/tfhers-ml/tfhers_params.json @@ -0,0 +1 @@ +{"lwe_dimension":902,"glwe_dimension":1,"polynomial_size":4096,"lwe_noise_distribution":{"Gaussian":{"std":1.0994794733558207e-6,"mean":0.0}},"glwe_noise_distribution":{"Gaussian":{"std":2.168404344971009e-19,"mean":0.0}},"pbs_base_log":15,"pbs_level":2,"ks_base_log":3,"ks_level":6,"message_modulus":4,"carry_modulus":8,"max_noise_level":10,"log2_p_fail":-64.084,"ciphertext_modulus":{"modulus":0,"scalar_bits":64},"encryption_key_choice":"Big"}