mirror of
https://github.com/zama-ai/concrete.git
synced 2026-04-17 03:00:54 -04:00
feat(mlir): implement mlir conversion of constant indexing
This commit is contained in:
@@ -7,7 +7,7 @@ import numpy
|
||||
import pytest
|
||||
|
||||
from concrete.common.compilation import CompilationConfiguration
|
||||
from concrete.common.data_types.integers import Integer, UnsignedInteger
|
||||
from concrete.common.data_types.integers import Integer, SignedInteger, UnsignedInteger
|
||||
from concrete.common.debugging import draw_graph, format_operation_graph
|
||||
from concrete.common.extensions.multi_table import MultiLookupTable
|
||||
from concrete.common.extensions.table import LookupTable
|
||||
@@ -1403,24 +1403,6 @@ return %2
|
||||
""".strip() # noqa: E501
|
||||
),
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[0],
|
||||
{"x": EncryptedTensor(Integer(3, is_signed=True), shape=(2, 2))},
|
||||
[(numpy.random.randint(-4, 2 ** 2, size=(2, 2)),) for i in range(10)],
|
||||
(
|
||||
"""
|
||||
|
||||
function you are trying to compile isn't supported for MLIR lowering
|
||||
|
||||
%0 = x # EncryptedTensor<int3, shape=(2, 2)>
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ only unsigned integer inputs are supported
|
||||
%1 = %0[0] # EncryptedTensor<int3, shape=(2,)>
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ indexing is not supported for the time being
|
||||
return %1
|
||||
|
||||
""".strip() # noqa: E501
|
||||
),
|
||||
),
|
||||
pytest.param(
|
||||
no_fuse_unhandled,
|
||||
{"x": EncryptedScalar(Integer(2, False)), "y": EncryptedScalar(Integer(2, False))},
|
||||
@@ -1539,6 +1521,52 @@ def test_fail_compile(function, parameters, inputset, match, default_compilation
|
||||
assert str(excinfo.value) == match, str(excinfo.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function,parameters,inputset,match",
|
||||
[
|
||||
pytest.param(
|
||||
lambda x: (x * 1.5)[0, 1],
|
||||
{"x": EncryptedTensor(SignedInteger(3), shape=(2, 2))},
|
||||
[(numpy.random.randint(-4, 3, size=(2, 2)),) for i in range(10)],
|
||||
(
|
||||
"""
|
||||
|
||||
function you are trying to compile isn't supported for MLIR lowering
|
||||
|
||||
%0 = x # EncryptedTensor<int3, shape=(2, 2)>
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ only unsigned integer inputs are supported
|
||||
%1 = 1.5 # ClearScalar<float64>
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ only integer constants are supported
|
||||
%2 = mul(%0, %1) # EncryptedTensor<float64, shape=(2, 2)>
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ only integer multiplication is supported
|
||||
%3 = %2[0, 1] # EncryptedScalar<float64>
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ only integer outputs are supported
|
||||
return %3
|
||||
|
||||
""".strip() # noqa: E501
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_fail_compile_while_fusing_is_disabled(
|
||||
function, parameters, inputset, match, default_compilation_configuration
|
||||
):
|
||||
"""Test compile_numpy_function without fusing and with failing inputs"""
|
||||
|
||||
configuration_to_use = deepcopy(default_compilation_configuration)
|
||||
configuration_to_use.enable_topological_optimizations = False
|
||||
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
compile_numpy_function(
|
||||
function,
|
||||
parameters,
|
||||
inputset,
|
||||
configuration_to_use,
|
||||
)
|
||||
|
||||
assert str(excinfo.value) == match, str(excinfo.value)
|
||||
|
||||
|
||||
def test_small_inputset_no_fail():
|
||||
"""Test function compile_numpy_function_into_op_graph with an unacceptably small inputset"""
|
||||
compile_numpy_function_into_op_graph_and_measure_bounds(
|
||||
|
||||
@@ -5,7 +5,10 @@ import pytest
|
||||
|
||||
from concrete.common.data_types import UnsignedInteger
|
||||
from concrete.common.values import EncryptedScalar, EncryptedTensor
|
||||
from concrete.numpy import compile_numpy_function_into_op_graph_and_measure_bounds
|
||||
from concrete.numpy import (
|
||||
compile_numpy_function,
|
||||
compile_numpy_function_into_op_graph_and_measure_bounds,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -595,3 +598,183 @@ def test_invalid_constant_indexing_with_numpy_values(
|
||||
except Exception as error:
|
||||
assert str(error) == expected_error_message
|
||||
raise
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function,parameters,inputset,test_input,expected_output",
|
||||
[
|
||||
pytest.param(
|
||||
lambda x: x[0],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(3), shape=(3,)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 3, size=(3,)),) for _ in range(10)],
|
||||
([4, 2, 6],),
|
||||
4,
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[-1],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(3), shape=(3,)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 3, size=(3,)),) for _ in range(10)],
|
||||
([4, 2, 6],),
|
||||
6,
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[:3],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(3), shape=(4,)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 3, size=(4,)),) for _ in range(10)],
|
||||
([4, 2, 6, 1],),
|
||||
[4, 2, 6],
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[2:],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(3), shape=(4,)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 3, size=(4,)),) for _ in range(10)],
|
||||
([4, 2, 6, 1],),
|
||||
[6, 1],
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[1:3],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(3), shape=(4,)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 3, size=(4,)),) for _ in range(10)],
|
||||
([4, 2, 6, 1],),
|
||||
[2, 6],
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[::2],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(3), shape=(4,)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 3, size=(4,)),) for _ in range(10)],
|
||||
([4, 2, 6, 1],),
|
||||
[4, 6],
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[::-1],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(3), shape=(4,)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 3, size=(4,)),) for _ in range(10)],
|
||||
([4, 2, 6, 1],),
|
||||
[1, 6, 2, 4],
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[1, 0],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(6), shape=(3, 2)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 6, size=(3, 2)),) for _ in range(10)],
|
||||
([[11, 12], [21, 22], [31, 32]],),
|
||||
21,
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[:, :],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(6), shape=(3, 2)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 6, size=(3, 2)),) for _ in range(10)],
|
||||
([[11, 12], [21, 22], [31, 32]],),
|
||||
[[11, 12], [21, 22], [31, 32]],
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[0, :],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(6), shape=(3, 2)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 6, size=(3, 2)),) for _ in range(10)],
|
||||
([[11, 12], [21, 22], [31, 32]],),
|
||||
[11, 12],
|
||||
marks=pytest.mark.xfail(strict=True),
|
||||
),
|
||||
pytest.param(
|
||||
lambda x: x[:, 0],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(6), shape=(3, 2)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 6, size=(3, 2)),) for _ in range(10)],
|
||||
([[11, 12], [21, 22], [31, 32]],),
|
||||
[11, 21, 31],
|
||||
marks=pytest.mark.xfail(strict=True),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_constant_indexing_run_correctness(
|
||||
function,
|
||||
parameters,
|
||||
inputset,
|
||||
test_input,
|
||||
expected_output,
|
||||
default_compilation_configuration,
|
||||
):
|
||||
"""Test correctness of results when running a compiled function with tensor operators"""
|
||||
circuit = compile_numpy_function(
|
||||
function,
|
||||
parameters,
|
||||
inputset,
|
||||
default_compilation_configuration,
|
||||
)
|
||||
|
||||
numpy_test_input = tuple(
|
||||
item if isinstance(item, int) else np.array(item, dtype=np.uint8) for item in test_input
|
||||
)
|
||||
|
||||
output = circuit.run(*numpy_test_input)
|
||||
expected = np.array(expected_output, dtype=np.uint8)
|
||||
|
||||
assert np.array_equal(
|
||||
output, expected
|
||||
), f"""
|
||||
|
||||
Actual Output
|
||||
=============
|
||||
{output}
|
||||
|
||||
Expected Output
|
||||
===============
|
||||
{expected}
|
||||
|
||||
"""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function,parameters,inputset,match",
|
||||
[
|
||||
pytest.param(
|
||||
lambda x: x[0:1],
|
||||
{
|
||||
"x": EncryptedTensor(UnsignedInteger(3), shape=(3,)),
|
||||
},
|
||||
[(np.random.randint(0, 2 ** 3, size=(3,)),) for _ in range(10)],
|
||||
(
|
||||
"Indexing of EncryptedTensor<uint3, shape=(3,)> with [0:1] "
|
||||
"cannot be converted to MLIR yet"
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_constant_indexing_failed_compilation(
|
||||
function,
|
||||
parameters,
|
||||
inputset,
|
||||
match,
|
||||
default_compilation_configuration,
|
||||
):
|
||||
"""Test compilation failures of compiled function with constant indexing"""
|
||||
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
compile_numpy_function(
|
||||
function,
|
||||
parameters,
|
||||
inputset,
|
||||
default_compilation_configuration,
|
||||
)
|
||||
|
||||
assert str(excinfo.value) == match, str(excinfo.value)
|
||||
|
||||
Reference in New Issue
Block a user