mirror of
https://github.com/zama-ai/concrete.git
synced 2026-01-25 04:37:58 -05:00
431 lines
11 KiB
Python
431 lines
11 KiB
Python
"""
|
|
Tests of everything related to multi-circuit.
|
|
"""
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from concrete import fhe
|
|
|
|
# pylint: disable=missing-class-docstring, missing-function-docstring, no-self-argument, unused-variable, no-member, unused-argument, function-redefined, expression-not-assigned
|
|
# same disables for ruff:
|
|
# ruff: noqa: N805, E501, F841, ARG002, F811, B015
|
|
|
|
|
|
def test_get_wrong_attribute():
|
|
"""
|
|
Test that getting the wrong attribute fails.
|
|
"""
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def square(x):
|
|
return x**2
|
|
|
|
with pytest.raises(AttributeError, match="No attribute nothing"):
|
|
a = Module.nothing
|
|
|
|
|
|
def test_empty_module():
|
|
"""
|
|
Test that defining a module without functions is an error.
|
|
"""
|
|
with pytest.raises(
|
|
RuntimeError, match="Tried to define an @fhe.module without any @fhe.function"
|
|
):
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
def square(x):
|
|
return x**2
|
|
|
|
|
|
def test_wrong_config(helpers):
|
|
"""
|
|
Test that defining a module with wrong configuration raises an error.
|
|
"""
|
|
|
|
with pytest.raises(RuntimeError) as excinfo:
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def add(x):
|
|
return x + 2
|
|
|
|
inputset = [np.random.randint(1, 20, size=()) for _ in range(100)]
|
|
module = Module.compile(
|
|
{"add": inputset},
|
|
composable=False,
|
|
)
|
|
|
|
assert str(excinfo.value) == ("Module can only be compiled with `composable` activated.")
|
|
|
|
|
|
def test_wrong_info():
|
|
"""
|
|
Test that defining a module with wrong information raises an error.
|
|
"""
|
|
|
|
with pytest.raises(ValueError) as excinfo:
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def add(x, y, z, w):
|
|
return x + y
|
|
|
|
assert str(excinfo.value) == (
|
|
"Encryption statuses of parameters 'y', 'z' and 'w' of function 'add' are not provided"
|
|
)
|
|
|
|
with pytest.raises(ValueError) as excinfo:
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted", "y": "encrypted", "z": "encrypted", "w": "encrypted"})
|
|
def add(x):
|
|
return x
|
|
|
|
assert str(excinfo.value) == (
|
|
"Encryption statuses of 'y', 'z' and 'w' are provided but they are not a parameter of function 'add'"
|
|
)
|
|
|
|
|
|
def test_wrong_inputset(helpers):
|
|
"""
|
|
Test that a wrong inputset raises an error.
|
|
"""
|
|
|
|
with pytest.raises(ValueError) as excinfo:
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted", "y": "encrypted"})
|
|
def add(x, y):
|
|
return x + y
|
|
|
|
inputset = [np.random.randint(1, 20, size=()) for _ in range(100)]
|
|
module = Module.compile(
|
|
{"add": inputset},
|
|
)
|
|
|
|
assert str(excinfo.value) == (
|
|
"Input #0 of your inputset is not well formed (expected a tuple of 2 values got a single value)"
|
|
)
|
|
|
|
with pytest.raises(RuntimeError) as excinfo:
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted", "y": "encrypted"})
|
|
def add(x, y):
|
|
return x + y
|
|
|
|
inputset = []
|
|
module = Module.compile(
|
|
{"add": inputset},
|
|
)
|
|
|
|
assert str(excinfo.value) == ("Compiling function 'add' without an inputset is not supported")
|
|
|
|
|
|
def test_call_clear_circuits():
|
|
"""
|
|
Test that calling clear functions works.
|
|
"""
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def square(x):
|
|
return x**2
|
|
|
|
@fhe.function({"x": "encrypted", "y": "encrypted"})
|
|
def add_sub(x, y):
|
|
return (x + y), (x - y)
|
|
|
|
@fhe.function({"x": "encrypted", "y": "encrypted"})
|
|
def mul(x, y):
|
|
return x * y
|
|
|
|
assert Module.square(2) == 4
|
|
assert Module.add_sub(2, 3) == (5, -1)
|
|
assert Module.mul(3, 4) == 12
|
|
|
|
|
|
def test_call_clear_circuits_wrong_kwargs():
|
|
"""
|
|
Test that calling clear functions works.
|
|
"""
|
|
|
|
with pytest.raises(RuntimeError) as excinfo:
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def square(x):
|
|
return x**2
|
|
|
|
Module.square(x=2) == 4
|
|
assert str(excinfo.value) == ("Calling function 'square' with kwargs is not supported")
|
|
|
|
|
|
def test_autorounder(helpers):
|
|
rounder1 = fhe.AutoRounder(target_msbs=5)
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def function1(x):
|
|
y = x + 1000
|
|
z = fhe.round_bit_pattern(y, lsbs_to_remove=rounder1)
|
|
return np.sqrt(z).astype(np.int64)
|
|
|
|
inputset1 = range(1000)
|
|
fhe.AutoRounder.adjust(Module.function1, inputset1)
|
|
assert rounder1.lsbs_to_remove == 6
|
|
|
|
module = Module.compile({"function1": inputset1}, auto_adjust_rounders=True)
|
|
|
|
|
|
def test_autotruncator(helpers):
|
|
truncator1 = fhe.AutoTruncator(target_msbs=5)
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def function1(x):
|
|
y = x + 1000
|
|
z = fhe.truncate_bit_pattern(y, lsbs_to_remove=truncator1)
|
|
return np.sqrt(z).astype(np.int64)
|
|
|
|
inputset1 = range(1000)
|
|
fhe.AutoTruncator.adjust(Module.function1, inputset1)
|
|
assert truncator1.lsbs_to_remove == 6
|
|
|
|
module = Module.compile(
|
|
{"function1": inputset1},
|
|
auto_adjust_truncators=True,
|
|
)
|
|
|
|
|
|
def test_compile():
|
|
"""
|
|
Test that compiling a module works.
|
|
"""
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def inc(x):
|
|
return x + 1
|
|
|
|
@fhe.function({"x": "encrypted"})
|
|
def dec(x):
|
|
return x - 1
|
|
|
|
inputset = [np.random.randint(1, 20, size=()) for _ in range(100)]
|
|
Module.compile({"inc": inputset, "dec": inputset}, verbose=True)
|
|
artifacts = fhe.ModuleDebugArtifacts()
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
module = Module.compile(
|
|
{"inc": inputset, "dec": inputset},
|
|
module_artifacts=artifacts,
|
|
verbose=True,
|
|
use_insecure_key_cache=True,
|
|
enable_unsafe_features=True,
|
|
insecure_key_cache_location=tmp,
|
|
)
|
|
|
|
assert module.mlir is not None
|
|
assert module.keys is not None
|
|
module.keygen()
|
|
module.cleanup()
|
|
assert set(artifacts.functions.keys()) == {"inc", "dec"}
|
|
|
|
|
|
def test_compiled_wrong_attribute():
|
|
"""
|
|
Test that getting unexisting attribute on module fails
|
|
"""
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def inc(x):
|
|
return x + 1
|
|
|
|
@fhe.function({"x": "encrypted"})
|
|
def dec(x):
|
|
return x - 1
|
|
|
|
inputset = [np.random.randint(1, 20, size=()) for _ in range(100)]
|
|
Module.compile({"inc": inputset, "dec": inputset}, verbose=True)
|
|
artifacts = fhe.ModuleDebugArtifacts()
|
|
|
|
module = Module.compile(
|
|
{"inc": inputset, "dec": inputset},
|
|
module_artifacts=artifacts,
|
|
verbose=True,
|
|
)
|
|
|
|
with pytest.raises(AttributeError, match="No attribute nothing"):
|
|
a = module.nothing
|
|
|
|
|
|
def test_compiled_clear_call():
|
|
"""
|
|
Test that cleartext execution works on compiled objects.
|
|
"""
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def inc(x):
|
|
return x + 1
|
|
|
|
@fhe.function({"x": "encrypted"})
|
|
def dec(x):
|
|
return x - 1
|
|
|
|
inputset = [np.random.randint(1, 20, size=()) for _ in range(100)]
|
|
module = Module.compile(
|
|
{"inc": inputset, "dec": inputset},
|
|
)
|
|
|
|
assert module.inc(5) == 6
|
|
assert module.dec(5) == 4
|
|
|
|
|
|
def test_compiled_simulation(helpers):
|
|
"""
|
|
Test that simulation works on compiled objects.
|
|
"""
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def inc(x):
|
|
return x + 1
|
|
|
|
@fhe.function({"x": "encrypted"})
|
|
def dec(x):
|
|
return x - 1
|
|
|
|
inputset = [np.random.randint(1, 20, size=()) for _ in range(100)]
|
|
|
|
module = Module.compile(
|
|
{"inc": inputset, "dec": inputset},
|
|
fhe_simulation=True,
|
|
)
|
|
|
|
assert module.keys is None
|
|
assert module.inc.simulate(5) == 6
|
|
assert module.dec.simulate(5) == 4
|
|
|
|
|
|
@pytest.mark.graphviz
|
|
def test_print(helpers):
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def inc(x):
|
|
return x + 1 % 20
|
|
|
|
inputset = list(range(20))
|
|
module = Module.compile(
|
|
{"inc": inputset},
|
|
)
|
|
helpers.check_str(
|
|
"""
|
|
%0 = x # EncryptedScalar<uint5> ∈ [0, 19]
|
|
%1 = 1 # ClearScalar<uint1> ∈ [1, 1]
|
|
%2 = add(%0, %1) # EncryptedScalar<uint5> ∈ [1, 20]
|
|
return %2
|
|
""",
|
|
str(module.inc),
|
|
)
|
|
|
|
|
|
def test_encrypted_execution():
|
|
"""
|
|
Test that encrypted execution works.
|
|
"""
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def inc(x):
|
|
return x + 1 % 20
|
|
|
|
@fhe.function({"x": "encrypted"})
|
|
def dec(x):
|
|
return x - 1 % 20
|
|
|
|
inputset = [np.random.randint(1, 20, size=()) for _ in range(100)]
|
|
module = Module.compile(
|
|
{"inc": inputset, "dec": inputset},
|
|
)
|
|
|
|
x = 5
|
|
x_enc = module.inc.encrypt(x)
|
|
x_inc_enc = module.inc.run(x_enc)
|
|
x_inc = module.inc.decrypt(x_inc_enc)
|
|
assert x_inc == 6
|
|
|
|
assert module.inc.encrypt_run_decrypt(2) == 3
|
|
|
|
x_inc_dec_enc = module.dec.run(x_inc_enc)
|
|
x_inc_dec = module.dec.decrypt(x_inc_dec_enc)
|
|
assert x_inc_dec == 5
|
|
|
|
for _ in range(10):
|
|
x_enc = module.inc.run(x_enc)
|
|
x_dec = module.inc.decrypt(x_enc)
|
|
assert x_dec == 15
|
|
|
|
|
|
def test_key_set():
|
|
"""
|
|
Test that keys can be set.
|
|
"""
|
|
|
|
@fhe.module()
|
|
class Module:
|
|
@fhe.function({"x": "encrypted"})
|
|
def inc(x):
|
|
return x + 1 % 20
|
|
|
|
@fhe.function({"x": "encrypted"})
|
|
def dec(x):
|
|
return x - 1 % 20
|
|
|
|
inputset = [np.random.randint(1, 20, size=()) for _ in range(100)]
|
|
module = Module.compile(
|
|
{"inc": inputset, "dec": inputset},
|
|
)
|
|
|
|
keys = module.keys
|
|
module.keygen(force=True)
|
|
module.keys = keys
|
|
|
|
x = 5
|
|
x_enc = module.inc.encrypt(x)
|
|
x_inc_enc = module.inc.run(x_enc)
|
|
x_inc = module.inc.decrypt(x_inc_enc)
|
|
assert x_inc == 6
|
|
|
|
assert module.inc.encrypt_run_decrypt(2) == 3
|
|
|
|
x_inc_dec_enc = module.dec.run(x_inc_enc)
|
|
x_inc_dec = module.dec.decrypt(x_inc_dec_enc)
|
|
assert x_inc_dec == 5
|
|
|
|
for _ in range(10):
|
|
x_enc = module.inc.run(x_enc)
|
|
x_dec = module.inc.decrypt(x_enc)
|
|
assert x_dec == 15
|