mirror of
https://github.com/zama-ai/concrete.git
synced 2026-01-13 06:48:02 -05:00
143 lines
4.5 KiB
Python
143 lines
4.5 KiB
Python
import os
|
|
from functools import partial
|
|
|
|
import click
|
|
import numpy as np
|
|
|
|
import concrete.fhe as 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 = False
|
|
#######################################
|
|
|
|
tfhers_type = tfhers.get_type_from_params(
|
|
TFHERS_PARAMS_FILE,
|
|
is_signed=IS_SIGNED,
|
|
precision=FHEUINT_PRECISION,
|
|
)
|
|
tfhers_int = partial(tfhers.TFHERSInteger, tfhers_type)
|
|
|
|
|
|
def compute(tfhers_x, tfhers_y):
|
|
####### 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)
|
|
concrete_y = tfhers.to_native(tfhers_y)
|
|
####### TFHE-rs to Concrete #########
|
|
|
|
####### Concrete Computation ########
|
|
concrete_res = (concrete_x + concrete_y) % 213
|
|
####### 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", "tfhers_y": "encrypted"})
|
|
|
|
inputset = [(tfhers_int(120), tfhers_int(120))]
|
|
circuit = compiler.compile(inputset)
|
|
|
|
tfhers_bridge = tfhers.new_bridge(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, 1: 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 Evaluation Keys")
|
|
with open(concrete_keyset_path, "wb") as f:
|
|
f.write(circuit.client.evaluation_keys.serialize())
|
|
print(f"saved Concrete Evaluation Keys 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("-c2", "--rust-ct-2", type=str, required=True)
|
|
@click.option("-o", "--output-rust-ct", type=str, required=True)
|
|
@click.option("-k", "--concrete-keyset-path", type=str, required=True)
|
|
def run(rust_ct_1: str, rust_ct_2: 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}'")
|
|
with open(concrete_keyset_path, "rb") as f:
|
|
eval_keys = fhe.EvaluationKeys.deserialize(f.read())
|
|
|
|
# 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)
|
|
|
|
# read tfhers int from file
|
|
with open(rust_ct_2, "rb") as f:
|
|
buff = f.read()
|
|
# import fheuint8 and get its description
|
|
tfhers_uint8_y = tfhers_bridge.import_value(buff, input_idx=1)
|
|
|
|
encrypted_x, encrypted_y = tfhers_uint8_x, tfhers_uint8_y
|
|
|
|
print("Homomorphic evaluation...")
|
|
encrypted_result = circuit.server.run(encrypted_x, encrypted_y, evaluation_keys=eval_keys)
|
|
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|