From 8f2b12d812b8ab2cf694c35346d1f4faf14527ec Mon Sep 17 00:00:00 2001 From: Antoniu Pop Date: Fri, 20 May 2022 12:53:44 +0100 Subject: [PATCH] test(distributed): fix termination of tests meant to run on distributed systems to avoid ungraceful exit. --- compiler/hpx.ini | 2 +- .../end_to_end_jit_auto_parallelization.cc | 58 ++++++++++--------- .../end_to_end_jit_clear_tensor.cc | 1 + .../end_to_end_jit_distributed.cc | 1 + .../end_to_end_jit_distributed.sh | 2 +- .../end_to_end_jit_encrypted_tensor.cc | 1 + .../end_to_end_tests/end_to_end_jit_fhe.cc | 1 + .../end_to_end_jit_fhelinalg.cc | 1 + .../end_to_end_tests/end_to_end_jit_lambda.cc | 1 + .../end_to_end_tests/end_to_end_jit_test.cc | 1 + compiler/tests/end_to_end_tests/globals.cc | 4 ++ compiler/tests/tests_tools/GtestEnvironment.h | 19 ++++++ .../TestLib/testlib_unit_test.cpp | 4 ++ 13 files changed, 67 insertions(+), 29 deletions(-) create mode 100644 compiler/tests/tests_tools/GtestEnvironment.h diff --git a/compiler/hpx.ini b/compiler/hpx.ini index 5427fe985..03605df68 100644 --- a/compiler/hpx.ini +++ b/compiler/hpx.ini @@ -3,7 +3,7 @@ location = ${HPX_LOCATION:$[system.prefix]} component_path = $[hpx.location]/lib/hpx:$[system.executable_prefix]/lib/hpx:$[system.executable_prefix]/../lib/hpx master_ini_path = $[hpx.location]/share/hpx-:$[system.executable_prefix]/share/hpx-:$[system.executable_prefix]/../share/hpx- ini_path = $[hpx.master_ini_path]/ini -os_threads = 1 +os_threads = 2 localities = 1 program_name = cmd_line = diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_auto_parallelization.cc b/compiler/tests/end_to_end_tests/end_to_end_jit_auto_parallelization.cc index 79e2eeac8..bd8eac2d6 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_auto_parallelization.cc +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_auto_parallelization.cc @@ -6,6 +6,7 @@ #include #include "end_to_end_jit_test.h" +#include "tests_tools/GtestEnvironment.h" /////////////////////////////////////////////////////////////////////////////// // Auto-parallelize independent FHE ops ///////////////////////////////////// @@ -73,10 +74,10 @@ func.func @main(%arg0: !FHE.eint<7>, %arg1: !FHE.eint<7>, %arg2: !FHE.eint<7>, % ASSERT_EXPECTED_VALUE(res_3, 60); ASSERT_EXPECTED_VALUE(res_4, 28); } else { - ASSERT_EXPECTED_FAILURE(lambda(1_u64, 2_u64, 3_u64, 4_u64)); - ASSERT_EXPECTED_FAILURE(lambda(4_u64, 5_u64, 6_u64, 7_u64)); - ASSERT_EXPECTED_FAILURE(lambda(1_u64, 1_u64, 1_u64, 1_u64)); - ASSERT_EXPECTED_FAILURE(lambda(5_u64, 7_u64, 11_u64, 13_u64)); + ASSERT_EXPECTED_FAILURE(lambda()); + ASSERT_EXPECTED_FAILURE(lambda()); + ASSERT_EXPECTED_FAILURE(lambda()); + ASSERT_EXPECTED_FAILURE(lambda()); } } @@ -119,12 +120,13 @@ TEST(ParallelizeAndRunFHE, nn_small_parallel) { ASSERT_EQ(res->size(), dim0 * dim2); parallel_results = *res; } else { - ASSERT_EXPECTED_FAILURE(lambda.operator()>({&arg})); + ASSERT_EXPECTED_FAILURE(lambda.operator()>()); } } TEST(ParallelizeAndRunFHE, nn_small_sequential) { - checkedJit(lambda, R"XXX( + if (mlir::concretelang::dfr::_dfr_is_root_node()) { + checkedJit(lambda, R"XXX( func @main(%arg0: tensor<4x5x!FHE.eint<5>>) -> tensor<4x7x!FHE.eint<5>> { %cst = arith.constant dense<[[0, 0, 1, 0, 1, 1, 0], [1, 1, 1, 0, 1, 0, 0], [1, 1, 0, 0, 0, 0, 0], [0, 0, 0, 0, 1, 1, 1]]> : tensor<4x7xi6> %cst_0 = arith.constant dense<[[1, 0, 1, 1, 0, 1, 1], [0, 1, 0, 0, 0, 0, 1], [0, 1, 1, 1, 1, 0, 0], [0, 1, 1, 0, 0, 0, 0], [0, 1, 1, 0, 0, 0, 1]]> : tensor<5x7xi6> @@ -135,30 +137,32 @@ TEST(ParallelizeAndRunFHE, nn_small_sequential) { return %2 : tensor<4x7x!FHE.eint<5>> } )XXX", - "main", false, false, false); + "main", false, false, false); - const size_t numDim = 2; - const size_t dim0 = 4; - const size_t dim1 = 5; - const size_t dim2 = 7; - const int64_t dims[numDim]{dim0, dim1}; - const llvm::ArrayRef shape2D(dims, numDim); - std::vector input; - input.reserve(dim0 * dim1); + const size_t numDim = 2; + const size_t dim0 = 4; + const size_t dim1 = 5; + const size_t dim2 = 7; + const int64_t dims[numDim]{dim0, dim1}; + const llvm::ArrayRef shape2D(dims, numDim); + std::vector input; + input.reserve(dim0 * dim1); - for (int i = 0; i < dim0 * dim1; ++i) - input.push_back(i % 17 % 4); + for (int i = 0; i < dim0 * dim1; ++i) + input.push_back(i % 17 % 4); - mlir::concretelang::TensorLambdaArgument< - mlir::concretelang::IntLambdaArgument> - arg(input, shape2D); + mlir::concretelang::TensorLambdaArgument< + mlir::concretelang::IntLambdaArgument> + arg(input, shape2D); - // This is sequential: only execute on root node. - if (mlir::concretelang::dfr::_dfr_is_root_node()) { - llvm::Expected> res = - lambda.operator()>({&arg}); - ASSERT_EXPECTED_SUCCESS(res); - for (size_t i = 0; i < dim0 * dim2; i++) - EXPECT_EQ(parallel_results[i], (*res)[i]) << "result differ at pos " << i; + // This is sequential: only execute on root node. + if (mlir::concretelang::dfr::_dfr_is_root_node()) { + llvm::Expected> res = + lambda.operator()>({&arg}); + ASSERT_EXPECTED_SUCCESS(res); + for (size_t i = 0; i < dim0 * dim2; i++) + EXPECT_EQ(parallel_results[i], (*res)[i]) + << "result differ at pos " << i; + } } } diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_clear_tensor.cc b/compiler/tests/end_to_end_tests/end_to_end_jit_clear_tensor.cc index 7ebf41ccb..c678603a7 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_clear_tensor.cc +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_clear_tensor.cc @@ -1,4 +1,5 @@ #include "end_to_end_jit_test.h" +#include "tests_tools/GtestEnvironment.h" /////////////////////////////////////////////////////////////////////////////// // 1D tensor ////////////////////////////////////////////////////////////////// diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.cc b/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.cc index b8f5821e7..adaa49d12 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.cc +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.cc @@ -5,6 +5,7 @@ #include #include "end_to_end_jit_test.h" +#include "tests_tools/GtestEnvironment.h" /////////////////////////////////////////////////////////////////////////////// // Auto-parallelize independent FHE ops ///////////////////////////////////// diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.sh b/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.sh index 3e4566bfa..58dfe57d1 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.sh +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.sh @@ -4,7 +4,7 @@ #SBATCH --mail-user=antoniu.pop@zama.ai #SBATCH --nodes=4 #SBATCH --cpus-per-task=8 -#SBATCH --time=00:20:00 +#SBATCH --time=00:45:00 #SBATCH --output=end_to_end_jit_distributed_%j.log echo "Date = $(date)" diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_encrypted_tensor.cc b/compiler/tests/end_to_end_tests/end_to_end_jit_encrypted_tensor.cc index aa4a7583f..ae9e52dcd 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_encrypted_tensor.cc +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_encrypted_tensor.cc @@ -1,4 +1,5 @@ #include "end_to_end_jit_test.h" +#include "tests_tools/GtestEnvironment.h" TEST(End2EndJit_EncryptedTensor_2D, extract_slice_parametric_2x2) { checkedJit(lambda, R"XXX( diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_fhe.cc b/compiler/tests/end_to_end_tests/end_to_end_jit_fhe.cc index e59fc68ca..533424cd1 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_fhe.cc +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_fhe.cc @@ -7,6 +7,7 @@ #include "concretelang/Support/LibrarySupport.h" #include "end_to_end_fixture/EndToEndFixture.h" #include "end_to_end_jit_test.h" +#include "tests_tools/GtestEnvironment.h" #include "tests_tools/keySetCache.h" template diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_fhelinalg.cc b/compiler/tests/end_to_end_tests/end_to_end_jit_fhelinalg.cc index 56ea878be..35e3642b2 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_fhelinalg.cc +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_fhelinalg.cc @@ -1,4 +1,5 @@ #include "end_to_end_jit_test.h" +#include "tests_tools/GtestEnvironment.h" namespace Z = mlir::concretelang; template diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_lambda.cc b/compiler/tests/end_to_end_tests/end_to_end_jit_lambda.cc index e463539d6..3d67db31b 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_lambda.cc +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_lambda.cc @@ -1,6 +1,7 @@ #include #include "end_to_end_jit_test.h" +#include "tests_tools/GtestEnvironment.h" TEST(Lambda_check_param, int_to_void_missing_param) { checkedJit(lambda, R"XXX( diff --git a/compiler/tests/end_to_end_tests/end_to_end_jit_test.cc b/compiler/tests/end_to_end_tests/end_to_end_jit_test.cc index ffaa21622..e0c36aa22 100644 --- a/compiler/tests/end_to_end_tests/end_to_end_jit_test.cc +++ b/compiler/tests/end_to_end_tests/end_to_end_jit_test.cc @@ -4,6 +4,7 @@ #include #include "end_to_end_jit_test.h" +#include "tests_tools/GtestEnvironment.h" TEST(CompileAndRunClear, add_u64) { checkedJit(lambda, R"XXX( diff --git a/compiler/tests/end_to_end_tests/globals.cc b/compiler/tests/end_to_end_tests/globals.cc index f42feaf59..d269f9da9 100644 --- a/compiler/tests/end_to_end_tests/globals.cc +++ b/compiler/tests/end_to_end_tests/globals.cc @@ -1,3 +1,7 @@ #include "globals.h" +#include "tests_tools/GtestEnvironment.h" const mlir::concretelang::V0FHEConstraint defaultV0Constraints{10, 7}; + +testing::Environment *const dfr_env = + testing::AddGlobalTestEnvironment(new DFREnvironment); diff --git a/compiler/tests/tests_tools/GtestEnvironment.h b/compiler/tests/tests_tools/GtestEnvironment.h new file mode 100644 index 000000000..ea47db552 --- /dev/null +++ b/compiler/tests/tests_tools/GtestEnvironment.h @@ -0,0 +1,19 @@ +#ifndef GTEST_ENVIRONMENT_H +#define GTEST_ENVIRONMENT_H + +#include "concretelang/Runtime/DFRuntime.hpp" +#include + +class DFREnvironment : public ::testing::Environment { +public: + ~DFREnvironment() override {} + + // Override this to define how to set up the environment. + void SetUp() override {} + + // Override this to define how to tear down the environment. + void TearDown() override { _dfr_terminate(); } +}; + +extern testing::Environment *const dfr_env; +#endif diff --git a/compiler/tests/unit_tests/concretelang/TestLib/testlib_unit_test.cpp b/compiler/tests/unit_tests/concretelang/TestLib/testlib_unit_test.cpp index ee0736269..98f836436 100644 --- a/compiler/tests/unit_tests/concretelang/TestLib/testlib_unit_test.cpp +++ b/compiler/tests/unit_tests/concretelang/TestLib/testlib_unit_test.cpp @@ -11,11 +11,15 @@ #include "concretelang/Support/CompilerEngine.h" #include "concretelang/TestLib/TestTypedLambda.h" +#include "tests_tools/GtestEnvironment.h" #include "tests_tools/assert.h" #include "tests_tools/keySetCache.h" #include "call_2t_1s_with_header-client.h.generated" +testing::Environment *const dfr_env = + testing::AddGlobalTestEnvironment(new DFREnvironment); + const std::string FUNCNAME = "main"; using namespace concretelang::testlib;