test(compiler): make benchmarks framework compatible with distributed…

… execution.
zama-ai · Dec 13, 2023 · 5f22831 · 5f22831
1 parent 30ae5f3
commit 5f22831
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 11 deletions.
diff --git a/compilers/concrete-compiler/compiler/tests/end_to_end_benchmarks/end_to_end_benchmark.cpp b/compilers/concrete-compiler/compiler/tests/end_to_end_benchmarks/end_to_end_benchmark.cpp
@@ -1,6 +1,7 @@
 #include "../end_to_end_tests/end_to_end_test.h"
 #include "concretelang/Common/Compat.h"
 #include "concretelang/TestLib/TestCircuit.h"
+#include <concretelang/Runtime/DFRuntime.hpp>
 
 #include <benchmark/benchmark.h>
 #include <filesystem>
@@ -19,6 +20,8 @@ using namespace concretelang::testlib;
     assert(false && "See error above");                                        \
   }
 
+static int num_iterations = 0;
+
 /// Benchmark time of the compilation
 static void BM_Compile(benchmark::State &state, EndToEndDesc description,
                        mlir::concretelang::CompilerEngine engine,
@@ -74,10 +77,12 @@ static void BM_ExportArguments(benchmark::State &state,
   auto inputArguments = std::vector<TransportValue>();
   inputArguments.reserve(test.inputs.size());
 
-  for (auto _ : state) {
-    for (size_t i = 0; i < test.inputs.size(); i++) {
-      auto input = circuit.prepareInput(test.inputs[i].getValue(), i).value();
-      inputArguments.push_back(input);
+  if (mlir::concretelang::dfr::_dfr_is_root_node()) {
+    for (auto _ : state) {
+      for (size_t i = 0; i < test.inputs.size(); i++) {
+        auto input = circuit.prepareInput(test.inputs[i].getValue(), i).value();
+        inputArguments.push_back(input);
+      }
     }
   }
 }
@@ -107,10 +112,12 @@ static void BM_Evaluate(benchmark::State &state, EndToEndDesc description,
   auto inputArguments = std::vector<TransportValue>();
   inputArguments.reserve(test.inputs.size());
 
-  for (size_t i = 0; i < test.inputs.size(); i++) {
-    auto input =
-        clientCircuit.prepareInput(test.inputs[i].getValue(), i).value();
-    inputArguments.push_back(input);
+  if (mlir::concretelang::dfr::_dfr_is_root_node()) {
+    for (size_t i = 0; i < test.inputs.size(); i++) {
+      auto input =
+          clientCircuit.prepareInput(test.inputs[i].getValue(), i).value();
+      inputArguments.push_back(input);
+    }
   }
 
   auto serverProgram = ServerProgram::load(
@@ -178,10 +185,12 @@ void registerEndToEndBenchmark(std::string suiteName,
             });
         break;
       case Action::EVALUATE:
-        benchmark::RegisterBenchmark(
+        auto bench = benchmark::RegisterBenchmark(
             benchName("evaluate").c_str(), [=](::benchmark::State &st) {
               BM_Evaluate(st, description, engine, options);
             });
+        if (num_iterations)
+          bench->Iterations(num_iterations);
         break;
       }
     }
@@ -210,6 +219,7 @@ int main(int argc, char **argv) {
 
   auto compilationOptions = std::get<0>(options);
   auto descriptionFiles = std::get<1>(options);
+  num_iterations = std::get<3>(options);
 
   std::vector<enum Action> actions = clActions;
   if (actions.empty()) {
@@ -226,5 +236,6 @@ int main(int argc, char **argv) {
   }
   ::benchmark::RunSpecifiedBenchmarks();
   ::benchmark::Shutdown();
+  _dfr_terminate();
   return 0;
 }
diff --git a/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.h b/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.h
@@ -23,7 +23,7 @@ const double TEST_ERROR_RATE = 1.0 - 0.999936657516;
 /// options, the library path if the --library options has been specified and
 /// the parsed description files
 std::tuple<mlir::concretelang::CompilationOptions,
-           std::vector<EndToEndDescFile>, int>
+           std::vector<EndToEndDescFile>, int, int>
 parseEndToEndCommandLine(int argc, char **argv) {
   namespace optimizer = mlir::concretelang::optimizer;
   // TODO - Well reset other llvm command line options registered but assert on
@@ -64,6 +64,18 @@ parseEndToEndCommandLine(int argc, char **argv) {
           "Set the batchTFHEOps compilation options to run the tests"),
       llvm::cl::init(false));
 
+  llvm::cl::opt<bool> distBenchmark(
+      "distributed",
+      llvm::cl::desc("Force a constant number of iterations in the benchmark "
+                     "suite as required for distributed execution (default: 1 "
+                     "- use --iterations=<n> to change)"),
+      llvm::cl::init(false));
+  llvm::cl::opt<int> numIterations(
+      "iterations",
+      llvm::cl::desc("Set the number of iterations for the benchmark suite "
+                     "(only to be used with --distributed)"),
+      llvm::cl::init(1));
+
   // Optimizer options
   llvm::cl::opt<int> securityLevel(
       "security-level",
@@ -129,8 +141,10 @@ parseEndToEndCommandLine(int argc, char **argv) {
     parsedDescriptionFiles.push_back(f);
   }
 
+  int num_iterations =
+      (distBenchmark.getValue()) ? numIterations.getValue() : 0;
   return std::make_tuple(compilationOptions, parsedDescriptionFiles,
-                         retryFailingTests.getValue());
+                         retryFailingTests.getValue(), num_iterations);
 }
 
 std::string getOptionsName(mlir::concretelang::CompilationOptions options) {