mirror of
https://github.com/zama-ai/concrete.git
synced 2026-02-10 04:35:03 -05:00
242 lines
8.2 KiB
C++
242 lines
8.2 KiB
C++
#include <gtest/gtest.h>
|
|
|
|
#include <numeric>
|
|
#include <cassert>
|
|
|
|
#include "../unittest/end_to_end_jit_test.h"
|
|
#include "concretelang/TestLib/DynamicLambda.h"
|
|
|
|
const std::string FUNCNAME = "main";
|
|
|
|
template<typename... Params>
|
|
using TypedDynamicLambda = mlir::concretelang::TypedDynamicLambda<Params...>;
|
|
|
|
using scalar = uint64_t;
|
|
using tensor1_in = std::vector<uint8_t>;
|
|
using tensor1_out = std::vector<uint64_t>;
|
|
using tensor2_out = std::vector<std::vector<uint64_t>>;
|
|
using tensor3_out = std::vector<std::vector<std::vector<uint64_t>>>;
|
|
|
|
std::vector<uint8_t>
|
|
values_7bits() {
|
|
return {0, 1, 2, 63, 64, 65, 125, 126};
|
|
}
|
|
|
|
llvm::Expected<mlir::concretelang::CompilerEngine::Library>
|
|
compile(std::string outputLib, std::string source) {
|
|
std::vector<std::string> sources = {source};
|
|
std::shared_ptr<mlir::concretelang::CompilationContext> ccx =
|
|
mlir::concretelang::CompilationContext::createShared();
|
|
mlir::concretelang::JitCompilerEngine ce {ccx};
|
|
ce.setClientParametersFuncName(FUNCNAME);
|
|
return ce.compile(sources, outputLib);
|
|
}
|
|
|
|
template<typename Info>
|
|
std::string outputLibFromThis(Info *info) {
|
|
return "tests/TestLib/out/" + std::string(info->name());
|
|
}
|
|
|
|
TEST(CompiledModule, call_1s_1s) {
|
|
std::string source = R"(
|
|
func @main(%arg0: !FHE.eint<7>) -> !FHE.eint<7> {
|
|
return %arg0: !FHE.eint<7>
|
|
}
|
|
)";
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<scalar, scalar>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
for(auto a: values_7bits()) {
|
|
auto res = lambda->call(a);
|
|
ASSERT_EXPECTED_VALUE(res, a);
|
|
}
|
|
}
|
|
|
|
TEST(CompiledModule, call_2s_1s) {
|
|
std::string source = R"(
|
|
func @main(%arg0: !FHE.eint<7>, %arg1: !FHE.eint<7>) -> !FHE.eint<7> {
|
|
%1 = "FHE.add_eint"(%arg0, %arg1): (!FHE.eint<7>, !FHE.eint<7>) -> (!FHE.eint<7>)
|
|
return %1: !FHE.eint<7>
|
|
}
|
|
)";
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<scalar, scalar, scalar>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
for(auto a: values_7bits()) for(auto b: values_7bits()) {
|
|
auto res = lambda->call(a, b);
|
|
ASSERT_EXPECTED_VALUE(res, a + b);
|
|
}
|
|
}
|
|
|
|
TEST(CompiledModule, call_1s_1t) {
|
|
std::string source = R"(
|
|
func @main(%arg0: !FHE.eint<7>) -> tensor<1x!FHE.eint<7>> {
|
|
%1 = tensor.from_elements %arg0 : tensor<1x!FHE.eint<7>>
|
|
return %1: tensor<1x!FHE.eint<7>>
|
|
}
|
|
)";
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<tensor1_out, scalar>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
for(auto a: values_7bits()) {
|
|
auto res = lambda->call(a);
|
|
ASSERT_EXPECTED_SUCCESS(res);
|
|
tensor1_out v = res.get();
|
|
EXPECT_EQ(v[0], a);
|
|
}
|
|
}
|
|
|
|
TEST(CompiledModule, call_2s_1t) {
|
|
std::string source = R"(
|
|
func @main(%arg0: !FHE.eint<7>, %arg1: !FHE.eint<7>) -> tensor<2x!FHE.eint<7>> {
|
|
%1 = tensor.from_elements %arg0, %arg1 : tensor<2x!FHE.eint<7>>
|
|
return %1: tensor<2x!FHE.eint<7>>
|
|
}
|
|
)";
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<tensor1_out, scalar, scalar>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
for(auto a : values_7bits()) {
|
|
auto res = lambda->call(a, a+1);
|
|
ASSERT_EXPECTED_SUCCESS(res);
|
|
tensor1_out v = res.get();
|
|
EXPECT_EQ((scalar)v[0], a);
|
|
EXPECT_EQ((scalar)v[1], a + 1u);
|
|
}
|
|
}
|
|
|
|
TEST(CompiledModule, call_1t_1s) {
|
|
std::string source = R"(
|
|
func @main(%arg0: tensor<1x!FHE.eint<7>>) -> !FHE.eint<7> {
|
|
%c0 = arith.constant 0 : index
|
|
%1 = tensor.extract %arg0[%c0] : tensor<1x!FHE.eint<7>>
|
|
return %1: !FHE.eint<7>
|
|
}
|
|
)";
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<scalar, tensor1_in>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
for(uint8_t a : values_7bits()) {
|
|
tensor1_in ta = {a};
|
|
auto res = lambda->call(ta);
|
|
ASSERT_EXPECTED_VALUE(res, a);
|
|
}
|
|
}
|
|
|
|
TEST(CompiledModule, call_1t_1t) {
|
|
std::string source = R"(
|
|
func @main(%arg0: tensor<3x!FHE.eint<7>>) -> tensor<3x!FHE.eint<7>> {
|
|
return %arg0: tensor<3x!FHE.eint<7>>
|
|
}
|
|
)";
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<tensor1_out, tensor1_in>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
tensor1_in ta = {1, 2, 3};
|
|
auto res = lambda->call(ta);
|
|
ASSERT_EXPECTED_SUCCESS(res);
|
|
tensor1_out v = res.get();
|
|
for(size_t i = 0; i < v.size(); i++) {
|
|
EXPECT_EQ(v[i], ta[i]);
|
|
}
|
|
}
|
|
|
|
TEST(CompiledModule, call_2t_1s) {
|
|
std::string source = R"(
|
|
func @main(%arg0: tensor<3x!FHE.eint<7>>, %arg1: tensor<3x!FHE.eint<7>>) -> !FHE.eint<7> {
|
|
%1 = "FHELinalg.add_eint"(%arg0, %arg1) : (tensor<3x!FHE.eint<7>>, tensor<3x!FHE.eint<7>>) -> tensor<3x!FHE.eint<7>>
|
|
%c1 = arith.constant 1 : i8
|
|
%2 = tensor.from_elements %c1, %c1, %c1 : tensor<3xi8>
|
|
%3 = "FHELinalg.dot_eint_int"(%1, %2) : (tensor<3x!FHE.eint<7>>, tensor<3xi8>) -> !FHE.eint<7>
|
|
return %3: !FHE.eint<7>
|
|
}
|
|
)";
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<scalar, tensor1_in, std::array<uint8_t, 3>>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
tensor1_in ta {1, 2, 3};
|
|
std::array<uint8_t, 3> tb {5, 7, 9};
|
|
auto res = lambda->call(ta, tb);
|
|
auto expected = std::accumulate(ta.begin(), ta.end(), 0u) +
|
|
std::accumulate(tb.begin(), tb.end(), 0u);
|
|
ASSERT_EXPECTED_VALUE(res, expected);
|
|
}
|
|
|
|
TEST(CompiledModule, call_1tr2_1tr2) {
|
|
std::string source = R"(
|
|
func @main(%arg0: tensor<2x3x!FHE.eint<7>>) -> tensor<2x3x!FHE.eint<7>> {
|
|
return %arg0: tensor<2x3x!FHE.eint<7>>
|
|
}
|
|
)";
|
|
using tensor2_in = std::array<std::array<uint8_t, 3>, 2>;
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<tensor2_out, tensor2_in>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
tensor2_in ta = {{
|
|
{1, 2, 3},
|
|
{4, 5, 6}
|
|
}};
|
|
auto res = lambda->call(ta);
|
|
ASSERT_EXPECTED_SUCCESS(res);
|
|
tensor2_out v = res.get();
|
|
for(size_t i = 0; i < v.size(); i++) {
|
|
for(size_t j = 0; j < v.size(); j++) {
|
|
EXPECT_EQ(v[i][j], ta[i][j]);
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
TEST(CompiledModule, call_1tr3_1tr3) {
|
|
std::string source = R"(
|
|
func @main(%arg0: tensor<2x3x1x!FHE.eint<7>>) -> tensor<2x3x1x!FHE.eint<7>> {
|
|
return %arg0: tensor<2x3x1x!FHE.eint<7>>
|
|
}
|
|
)";
|
|
using tensor3_in = std::array<std::array<std::array<uint8_t, 1>, 3>, 2>;
|
|
std::string outputLib = outputLibFromThis(this->test_info_);
|
|
auto compiled = compile(outputLib, source);
|
|
ASSERT_EXPECTED_SUCCESS(compiled);
|
|
auto lambda = TypedDynamicLambda<tensor3_out, tensor3_in>::load(FUNCNAME, outputLib);
|
|
ASSERT_EXPECTED_SUCCESS(lambda);
|
|
ASSERT_LLVM_ERROR(lambda->generateKeySet(getTestKeySetCache()));
|
|
tensor3_in ta = {{
|
|
{{ {1}, {2}, {3} }},
|
|
{{ {4}, {5}, {6} }}
|
|
}};
|
|
auto res = lambda->call(ta);
|
|
ASSERT_EXPECTED_SUCCESS(res);
|
|
tensor3_out v = res.get();
|
|
for(size_t i = 0; i < v.size(); i++) {
|
|
for(size_t j = 0; j < v[i].size(); j++) {
|
|
for(size_t k = 0; k < v[i][j].size(); k++) {
|
|
EXPECT_EQ(v[i][j][k], ta[i][j][k]);
|
|
}
|
|
}
|
|
}
|
|
}
|