diff --git a/compiler/Makefile b/compiler/Makefile index 88a8763a7..ee84ac765 100644 --- a/compiler/Makefile +++ b/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 diff --git a/compiler/tests/end_to_end_benchmarks/CMakeLists.txt b/compiler/tests/end_to_end_benchmarks/CMakeLists.txt index 90209fa60..e1f0f7a3e 100644 --- a/compiler/tests/end_to_end_benchmarks/CMakeLists.txt +++ b/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") + diff --git a/compiler/tests/end_to_end_benchmarks/end_to_end_mlbench.cpp b/compiler/tests/end_to_end_benchmarks/end_to_end_mlbench.cpp new file mode 100644 index 000000000..36876213b --- /dev/null +++ b/compiler/tests/end_to_end_benchmarks/end_to_end_mlbench.cpp @@ -0,0 +1,180 @@ +#include "end_to_end_fixture/EndToEndFixture.h" +#include +#define BENCHMARK_HAS_CXX11 +#include "llvm/Support/Path.h" +#include + +#include "tests_tools/StackSize.h" +#include "tests_tools/keySetCache.h" + +/// Benchmark time of the compilation +template +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 +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 +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 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 +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 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; +} diff --git a/compiler/tests/end_to_end_benchmarks/end_to_end_mlbench.sh b/compiler/tests/end_to_end_benchmarks/end_to_end_mlbench.sh new file mode 100755 index 000000000..7f2830f7a --- /dev/null +++ b/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" \; diff --git a/compiler/tests/end_to_end_benchmarks/generate_bench_yaml.py b/compiler/tests/end_to_end_benchmarks/generate_bench_yaml.py new file mode 100644 index 000000000..b509d0613 --- /dev/null +++ b/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') + diff --git a/compiler/tests/end_to_end_benchmarks/mlbench.zip b/compiler/tests/end_to_end_benchmarks/mlbench.zip new file mode 100644 index 000000000..e629806ba Binary files /dev/null and b/compiler/tests/end_to_end_benchmarks/mlbench.zip differ