From 39e7313348beea9a16a0db41ea0a0cf7a4c40091 Mon Sep 17 00:00:00 2001
From: Antoniu Pop <antoniu.pop@zama.ai>
Date: Thu, 1 Sep 2022 13:02:49 +0100
Subject: [PATCH] feat(tests): add function for increasing the stack limit and
 use for benchmark overflow.

---
 .../end_to_end_benchmark.cpp                  |  8 ++++--
 compiler/tests/tests_tools/StackSize.h        | 26 +++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100644 compiler/tests/tests_tools/StackSize.h

diff --git a/compiler/tests/end_to_end_benchmarks/end_to_end_benchmark.cpp b/compiler/tests/end_to_end_benchmarks/end_to_end_benchmark.cpp
index 22f1705ea..9fc21f5de 100644
--- a/compiler/tests/end_to_end_benchmarks/end_to_end_benchmark.cpp
+++ b/compiler/tests/end_to_end_benchmarks/end_to_end_benchmark.cpp
@@ -3,6 +3,7 @@
 #include "llvm/Support/Path.h"
 #include <benchmark/benchmark.h>
 
+#include "tests_tools/StackSize.h"
 #include "tests_tools/keySetCache.h"
 
 /// Benchmark time of the compilation
@@ -100,7 +101,8 @@ static void BM_Evaluate(benchmark::State &state, EndToEndDesc description,
   }
 }
 
-static int registerEndToEndTestFromFile(std::string prefix, std::string path) {
+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) {
@@ -131,6 +133,7 @@ static int registerEndToEndTestFromFile(std::string prefix, std::string path) {
       return;
     });
   };
+  setCurrentStackLimit(stackSizeRequirement);
   mlir::concretelang::CompilationOptions defaul;
   registe("default", defaul);
   mlir::concretelang::CompilationOptions loop;
@@ -152,6 +155,7 @@ auto _ = {
     registerEndToEndTestFromFile(
         "FHELinalg", "tests/end_to_end_fixture/end_to_end_fhelinalg.yaml"),
     registerEndToEndTestFromFile(
-        "FHELinalg", "tests/end_to_end_fixture/end_to_end_programs.yaml")};
+        "FHELinalg", "tests/end_to_end_fixture/end_to_end_programs.yaml",
+        0x8000000)};
 
 BENCHMARK_MAIN();
diff --git a/compiler/tests/tests_tools/StackSize.h b/compiler/tests/tests_tools/StackSize.h
new file mode 100644
index 000000000..8c69aaade
--- /dev/null
+++ b/compiler/tests/tests_tools/StackSize.h
@@ -0,0 +1,26 @@
+#ifndef TEST_TOOLS_STACK_SIZE_H
+#define TEST_TOOLS_STACK_SIZE_H
+
+#include <iostream>
+#include <sys/resource.h>
+
+void setCurrentStackLimit(size_t size) {
+  rlim_t stack_size_target = size;
+  struct rlimit rl;
+
+  int res = getrlimit(RLIMIT_STACK, &rl);
+  if (!res) {
+    if (rl.rlim_cur < stack_size_target) {
+      rl.rlim_cur = stack_size_target;
+      res = setrlimit(RLIMIT_STACK, &rl);
+      if (res)
+        std::cerr << "Unable to set the requiresd stack size ("
+                  << stack_size_target << " Bytes) - setrlimit returned " << res
+                  << std::endl;
+      else
+        std::cerr << "Stack limit increased to " << rl.rlim_cur << " Bytes\n";
+    }
+  }
+}
+
+#endif