feat(mlir): implement mlir conversion of constant indexing

This commit is contained in:
Umut
2021-11-29 10:23:05 +03:00
parent 926597c3f6
commit ac74e94e13
4 changed files with 335 additions and 22 deletions

View File

@@ -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(

View File

@@ -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)