feat: support more dtype for scalars/tensors

dtype supported now: uint8, uint16, uint32, uint64
This commit is contained in:
youben11
2021-12-10 10:36:04 +01:00
committed by Ayoub Benaissa
parent 550318f67e
commit 60b2cfd9b7
5 changed files with 143 additions and 12 deletions

View File

@@ -43,6 +43,39 @@ KEY_SET_CACHE_PATH = os.path.join(tempfile.gettempdir(), 'KeySetCache')
8,
id="add_eint_int_with_np_uint8_as_scalar",
),
pytest.param(
"""
func @main(%arg0: !HLFHE.eint<7>, %arg1: i8) -> !HLFHE.eint<7> {
%1 = "HLFHE.add_eint_int"(%arg0, %arg1): (!HLFHE.eint<7>, i8) -> (!HLFHE.eint<7>)
return %1: !HLFHE.eint<7>
}
""",
(np.uint16(3), np.uint16(5)),
8,
id="add_eint_int_with_np_uint16_as_scalar",
),
pytest.param(
"""
func @main(%arg0: !HLFHE.eint<7>, %arg1: i8) -> !HLFHE.eint<7> {
%1 = "HLFHE.add_eint_int"(%arg0, %arg1): (!HLFHE.eint<7>, i8) -> (!HLFHE.eint<7>)
return %1: !HLFHE.eint<7>
}
""",
(np.uint32(3), np.uint32(5)),
8,
id="add_eint_int_with_np_uint32_as_scalar",
),
pytest.param(
"""
func @main(%arg0: !HLFHE.eint<7>, %arg1: i8) -> !HLFHE.eint<7> {
%1 = "HLFHE.add_eint_int"(%arg0, %arg1): (!HLFHE.eint<7>, i8) -> (!HLFHE.eint<7>)
return %1: !HLFHE.eint<7>
}
""",
(np.uint64(3), np.uint64(5)),
8,
id="add_eint_int_with_np_uint64_as_scalar",
),
pytest.param(
"""
func @main(%arg0: tensor<4x!HLFHE.eint<7>>, %arg1: tensor<4xi8>) -> !HLFHE.eint<7>
@@ -57,7 +90,55 @@ KEY_SET_CACHE_PATH = os.path.join(tempfile.gettempdir(), 'KeySetCache')
np.array([4, 3, 2, 1], dtype=np.uint8),
),
20,
id="dot_eint_int",
id="dot_eint_int_uint8",
),
pytest.param(
"""
func @main(%arg0: tensor<4x!HLFHE.eint<7>>, %arg1: tensor<4xi8>) -> !HLFHE.eint<7>
{
%ret = "HLFHELinalg.dot_eint_int"(%arg0, %arg1) :
(tensor<4x!HLFHE.eint<7>>, tensor<4xi8>) -> !HLFHE.eint<7>
return %ret : !HLFHE.eint<7>
}
""",
(
np.array([1, 2, 3, 4], dtype=np.uint16),
np.array([4, 3, 2, 1], dtype=np.uint16),
),
20,
id="dot_eint_int_uint16",
),
pytest.param(
"""
func @main(%arg0: tensor<4x!HLFHE.eint<7>>, %arg1: tensor<4xi8>) -> !HLFHE.eint<7>
{
%ret = "HLFHELinalg.dot_eint_int"(%arg0, %arg1) :
(tensor<4x!HLFHE.eint<7>>, tensor<4xi8>) -> !HLFHE.eint<7>
return %ret : !HLFHE.eint<7>
}
""",
(
np.array([1, 2, 3, 4], dtype=np.uint32),
np.array([4, 3, 2, 1], dtype=np.uint32),
),
20,
id="dot_eint_int_uint32",
),
pytest.param(
"""
func @main(%arg0: tensor<4x!HLFHE.eint<7>>, %arg1: tensor<4xi8>) -> !HLFHE.eint<7>
{
%ret = "HLFHELinalg.dot_eint_int"(%arg0, %arg1) :
(tensor<4x!HLFHE.eint<7>>, tensor<4xi8>) -> !HLFHE.eint<7>
return %ret : !HLFHE.eint<7>
}
""",
(
np.array([1, 2, 3, 4], dtype=np.uint64),
np.array([4, 3, 2, 1], dtype=np.uint64),
),
20,
id="dot_eint_int_uint64",
),
pytest.param(
"""