benchmarking: add YAML generator and ML benchmarks (see issue https://github.com/zama-ai/concrete-ml-internal/issues/1543).

compiler/Makefile

@@ -58,7 +58,7 @@ else
 		PYTHON_TESTS_MARKER="not parallel"
 endif
 
-all: concretecompiler python-bindings build-tests build-benchmarks doc
+all: concretecompiler python-bindings build-tests build-benchmarks build-mlbench doc
 
 # concrete-core-ffi #######################################
 
@@ -233,6 +233,20 @@ build-benchmarks: build-initialized
 run-benchmarks: build-benchmarks
 	$(BUILD_DIR)/bin/end_to_end_benchmark
 
+build-mlbench: build-initialized
+	cmake --build $(BUILD_DIR) --target end_to_end_mlbench
+
+generate-mlbench:
+	mkdir -p tests/end_to_end_benchmarks/mlbench
+	rm -rf tests/end_to_end_benchmarks/mlbench/*
+	unzip tests/end_to_end_benchmarks/mlbench.zip -d tests/end_to_end_benchmarks/mlbench
+	rm -f tests/end_to_end_benchmarks/mlbench/**/*\=*
+	find tests/end_to_end_benchmarks/mlbench -name "*.mlir" -exec sed -e '1d' -e 's/ func / func.func /g' -e 's/ linalg.tensor_/ tensor./g' -e '$$d' -i {} \;
+	python tests/end_to_end_benchmarks/generate_bench_yaml.py tests/end_to_end_benchmarks/mlbench tests/end_to_end_benchmarks/mlbench/end_to_end_mlbench
+
+run-mlbench: build-mlbench
+	tests/end_to_end_benchmarks/end_to_end_mlbench.sh tests/end_to_end_benchmarks/mlbench/ $(BUILD_DIR)/bin/end_to_end_mlbench
+
 show-stress-tests-summary:
 	@echo '------ Stress tests summary ------'
 	@echo

compiler/tests/end_to_end_benchmarks/CMakeLists.txt

@@ -3,3 +3,7 @@ add_executable(end_to_end_benchmark end_to_end_benchmark.cpp)
 target_link_libraries(end_to_end_benchmark benchmark::benchmark ConcretelangSupport EndToEndFixture)
 set_source_files_properties(end_to_end_benchmark.cpp PROPERTIES COMPILE_FLAGS "-fno-rtti")
 
+add_executable(end_to_end_mlbench end_to_end_mlbench.cpp)
+target_link_libraries(end_to_end_mlbench benchmark::benchmark ConcretelangSupport EndToEndFixture)
+set_source_files_properties(end_to_end_mlbench.cpp PROPERTIES COMPILE_FLAGS "-fno-rtti")
+

compiler/tests/end_to_end_benchmarks/end_to_end_mlbench.cpp

@@ -0,0 +1,180 @@
+#include "end_to_end_fixture/EndToEndFixture.h"
+#include <concretelang/Runtime/DFRuntime.hpp>
+#define BENCHMARK_HAS_CXX11
+#include "llvm/Support/Path.h"
+#include <benchmark/benchmark.h>
+
+#include "tests_tools/StackSize.h"
+#include "tests_tools/keySetCache.h"
+
+/// Benchmark time of the compilation
+template <typename LambdaSupport>
+static void BM_Compile(benchmark::State &state, EndToEndDesc description,
+                       LambdaSupport support,
+                       mlir::concretelang::CompilationOptions options) {
+  for (auto _ : state) {
+    if (support.compile(description.program, options)) {
+    };
+  }
+}
+
+/// Benchmark time of the key generation
+template <typename LambdaSupport>
+static void BM_KeyGen(benchmark::State &state, EndToEndDesc description,
+                      LambdaSupport support,
+                      mlir::concretelang::CompilationOptions options) {
+  auto compilationResult = support.compile(description.program, options);
+  assert(compilationResult);
+
+  auto clientParameters = support.loadClientParameters(**compilationResult);
+  assert(clientParameters);
+
+  for (auto _ : state) {
+    assert(support.keySet(*clientParameters, llvm::None));
+  }
+}
+
+/// Benchmark time of the encryption
+template <typename LambdaSupport>
+static void BM_ExportArguments(benchmark::State &state,
+                               EndToEndDesc description, LambdaSupport support,
+                               mlir::concretelang::CompilationOptions options) {
+  auto compilationResult = support.compile(description.program, options);
+  assert(compilationResult);
+
+  auto clientParameters = support.loadClientParameters(**compilationResult);
+  assert(clientParameters);
+
+  auto keySet = support.keySet(*clientParameters, getTestKeySetCache());
+  assert(keySet);
+
+  assert(description.tests.size() > 0);
+  auto test = description.tests[0];
+  std::vector<mlir::concretelang::LambdaArgument *> inputArguments;
+  inputArguments.reserve(test.inputs.size());
+  for (auto input : test.inputs) {
+    auto arg = valueDescriptionToLambdaArgument(input);
+    assert(arg);
+    inputArguments.push_back(arg.get());
+  }
+
+  for (auto _ : state) {
+    assert(
+        support.exportArguments(*clientParameters, **keySet, inputArguments));
+  }
+}
+
+/// Benchmark time of the program evaluation
+template <typename LambdaSupport>
+static void BM_Evaluate(benchmark::State &state, EndToEndDesc description,
+                        LambdaSupport support,
+                        mlir::concretelang::CompilationOptions options) {
+  auto compilationResult = support.compile(description.program, options);
+  assert(compilationResult);
+
+  auto clientParameters = support.loadClientParameters(**compilationResult);
+  assert(clientParameters);
+
+  auto keySet = support.keySet(*clientParameters, getTestKeySetCache());
+  assert(keySet);
+
+  assert(description.tests.size() > 0);
+  auto test = description.tests[0];
+  std::vector<mlir::concretelang::LambdaArgument *> inputArguments;
+  inputArguments.reserve(test.inputs.size());
+  for (auto input : test.inputs) {
+    auto arg = valueDescriptionToLambdaArgument(input);
+    assert(arg);
+    inputArguments.push_back(arg.get());
+  }
+
+  auto publicArguments =
+      support.exportArguments(*clientParameters, **keySet, inputArguments);
+  assert(publicArguments);
+
+  auto serverLambda = support.loadServerLambda(**compilationResult);
+  assert(serverLambda);
+  auto evaluationKeys = (*keySet)->evaluationKeys();
+
+  for (auto _ : state) {
+    assert(
+        support.serverCall(*serverLambda, **publicArguments, evaluationKeys));
+  }
+}
+
+static int registerEndToEndTestFromFile(std::string prefix, std::string path,
+                                        size_t stackSizeRequirement = 0) {
+  auto registe = [&](std::string optionsName,
+                     mlir::concretelang::CompilationOptions options) {
+    llvm::for_each(loadEndToEndDesc(path), [&](EndToEndDesc &description) {
+      options.clientParametersFuncName = "main";
+      mlir::concretelang::JITSupport support;
+      auto benchName = [&](std::string name) {
+        std::ostringstream s;
+        s << prefix << "/" << name << "/" << optionsName << "/"
+          << description.description;
+        return s.str();
+      };
+      benchmark::RegisterBenchmark(
+          benchName("Compile").c_str(), [=](::benchmark::State &st) {
+            BM_Compile(st, description, support, options);
+          });
+      benchmark::RegisterBenchmark(
+          benchName("KeyGen").c_str(), [=](::benchmark::State &st) {
+            BM_KeyGen(st, description, support, options);
+          });
+      benchmark::RegisterBenchmark(
+          benchName("ExportArguments").c_str(), [=](::benchmark::State &st) {
+            BM_ExportArguments(st, description, support, options);
+          });
+      benchmark::RegisterBenchmark(
+          benchName("Evaluate").c_str(), [=](::benchmark::State &st) {
+            BM_Evaluate(st, description, support, options);
+          });
+      return;
+    });
+  };
+  setCurrentStackLimit(stackSizeRequirement);
+  mlir::concretelang::CompilationOptions defaul;
+  registe("default", defaul);
+  mlir::concretelang::CompilationOptions loop;
+  loop.loopParallelize = true;
+  registe("loop", loop);
+#ifdef CONCRETELANG_DATAFLOW_EXECUTION_ENABLED
+  mlir::concretelang::CompilationOptions dataflow;
+  dataflow.dataflowParallelize = true;
+  dataflow.loopParallelize = true;
+  registe("dataflow", dataflow);
+#endif
+  return 1;
+}
+
+int main(int argc, char **argv) {
+  char *bench_name = "MLBench";
+  char *file_name =
+      "tests/end_to_end_benchmarks/mlbench/end_to_end_mlbench.yaml";
+  size_t stack_size = 0;
+
+  char *env = getenv("BENCHMARK_NAME");
+  if (env != nullptr)
+    bench_name = env;
+  env = getenv("BENCHMARK_FILE");
+  if (env != nullptr)
+    file_name = env;
+  env = getenv("BENCHMARK_STACK");
+  if (env != nullptr)
+    stack_size = strtoul(env, NULL, 10);
+
+  std::cout << "Benchmark executing [" << bench_name << "] from file "
+            << file_name << "\n";
+  registerEndToEndTestFromFile(bench_name, file_name, stack_size);
+
+  ::benchmark::Initialize(&argc, argv);
+  if (::benchmark::ReportUnrecognizedArguments(argc, argv))
+    return 1;
+  ::benchmark::RunSpecifiedBenchmarks();
+  ::benchmark::Shutdown();
+
+  _dfr_terminate();
+  return 0;
+}

compiler/tests/end_to_end_benchmarks/end_to_end_mlbench.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+find $1 -name "*mlbench_*.yaml" -exec bash -c "BENCHMARK_FILE={} BENCHMARK_STACK=1000000000 BENCHMARK_NAME=MLBench $2" \;

compiler/tests/end_to_end_benchmarks/generate_bench_yaml.py

@@ -0,0 +1,85 @@
+import os
+import re
+import sys
+import numpy
+import random
+
+class Param:
+    def __init__(self):
+        self.typ = ""
+        self.values = []
+        self.shapes = []
+        self.width = 0
+
+def generateHeader(output, name):
+    output.write("description: " + name + "\n")
+    output.write("program: |\n")
+
+def generateFooter(output, params):
+    output.write("tests:\n")
+    output.write("  - inputs:\n")
+    for p in params:
+        if p.typ == "scalar":
+            output.write("    - scalar: " + p.value + "\n")
+        if p.typ == "tensor":
+            output.write("    - tensor: [")
+            for i, v in enumerate(p.values):
+                sv = str(v)
+                if i == 0:
+                    output.write(sv)
+                else:
+                    output.write(", " + sv)
+            output.write("]\n")
+            output.write("      shape: [")
+            for i, v in enumerate(p.shapes):
+                sv = str(v)
+                if i == 0:
+                    output.write(sv)
+                else:
+                    output.write(", " + sv)
+            output.write("]\n")
+            #output.write("      width: " + str(p.width+1) + "\n")
+    output.write("---\n\n")
+
+def getParams(filename):
+    f = open(filename, 'r')
+    params = []
+    for line in f:
+        m = re.match(r".*?func.func @main\((.*?)\).*?", line)
+        if m:
+            args = re.split(r'%\w+:', m.group(1))
+            for a in args:
+                am = re.match(r"\W*tensor<((?:\d+x)+)(?:(?:!FHE.eint<(\d+)>>)|(?:i(\d+)>))", a)
+                if am:
+                    param = Param()
+                    param.typ = "tensor"
+                    shapes = list(filter(None, re.split(r'x', am.group(1))))
+                    param.shapes = list(map(int, shapes))
+                    if am.group(2):
+                        param.width = int(am.group(2))
+                    else:
+                        param.width = int(am.group(3))
+                    for i in range(0, numpy.prod(param.shapes)):
+                        param.values.append(random.randint(0, 2**param.width))
+                    params.append(param)
+    return params
+
+if __name__ == "__main__":
+    # Find all MLIRs
+    for dirname, dirnames, filenames in os.walk(sys.argv[1]):
+        for i, filename in enumerate(filenames):
+            if i % 20 == 0:
+                output = open(sys.argv[2] + "_" + str(int(i/20)) + ".yaml", 'w')
+            desc = re.match(r"(.*?)\.mlir$", filename)
+            if desc:
+                generateHeader(output, desc.group(1))
+                f = open(os.path.join(dirname, filename), 'r')
+                output.write(f.read() + "\n")
+                f.close()
+                generateFooter(output, getParams(os.path.join(dirname, filename)))
+            if i % 20 == 19:
+                output.close()
+
+        if '.git' in dirnames:
+            dirnames.remove('.git')
+