fix(runtime/DFR): create only one BSK clone per thread and deallocate.

compiler/include/concretelang/ClientLib/KeySet.h

@@ -10,8 +10,8 @@
 
 extern "C" {
 #include "concrete-ffi.h"
-#include "concretelang/Runtime/context.h"
 }
+#include "concretelang/Runtime/context.h"
 
 #include "concretelang/ClientLib/ClientParameters.h"
 #include "concretelang/ClientLib/KeySetCache.h"
@@ -57,7 +57,8 @@ public:
 
   void setRuntimeContext(RuntimeContext &context) {
     context.ksk = std::get<1>(this->keyswitchKeys["ksk_v0"]);
-    context.bsk = std::get<1>(this->bootstrapKeys["bsk_v0"]);
+    context.bsk["_concretelang_base_context_bsk"] =
+        std::get<1>(this->bootstrapKeys["bsk_v0"]);
   }
 
   const std::map<LweSecretKeyID,
@@ -109,4 +110,4 @@ private:
 } // namespace concretelang
 } // namespace mlir
 
-#endif
+#endif

compiler/include/concretelang/Runtime/context.h

@@ -6,15 +6,21 @@
 #ifndef CONCRETELANG_RUNTIME_CONTEXT_H
 #define CONCRETELANG_RUNTIME_CONTEXT_H
 
+#include <map>
+#include <string>
+
+extern "C" {
 #include "concrete-ffi.h"
+}
 
 typedef struct RuntimeContext {
   LweKeyswitchKey_u64 *ksk;
-  LweBootstrapKey_u64 *bsk;
+  std::map<std::string, LweBootstrapKey_u64 *> bsk;
 } RuntimeContext;
 
+extern "C" {
 LweKeyswitchKey_u64 *get_keyswitch_key(RuntimeContext *context);
 
 LweBootstrapKey_u64 *get_bootstrap_key(RuntimeContext *context);
-
+}
 #endif

compiler/lib/Runtime/CMakeLists.txt

@@ -1,17 +1,23 @@
+if(CONCRETELANG_PARALLEL_EXECUTION_ENABLED)
+  add_compile_options(-DCONCRETELANG_PARALLEL_EXECUTION_ENABLED)
+endif()
+
 add_library(ConcretelangRuntime SHARED
-    context.c
+    context.cpp
     wrappers.c
 )
 
-target_link_libraries(ConcretelangRuntime Concrete pthread m dl)
-
-install(TARGETS ConcretelangRuntime EXPORT ConcretelangRuntime)
-install(EXPORT ConcretelangRuntime DESTINATION "./")
-
 if(CONCRETELANG_PARALLEL_EXECUTION_ENABLED)
   add_library(DFRuntime SHARED DFRuntime.cpp)
   target_link_libraries(DFRuntime PUBLIC pthread m dl HPX::hpx HPX::iostreams_component -rdynamic)
 
   install(TARGETS DFRuntime EXPORT DFRuntime)
   install(EXPORT DFRuntime DESTINATION "./")
+
+  target_link_libraries(ConcretelangRuntime Concrete pthread m dl HPX::hpx)
+else()
+  target_link_libraries(ConcretelangRuntime Concrete pthread m dl)
 endif()
+
+install(TARGETS ConcretelangRuntime EXPORT ConcretelangRuntime)
+install(EXPORT ConcretelangRuntime DESTINATION "./")

compiler/lib/Runtime/context.c

@@ -1,16 +0,0 @@
-#include "concrete-ffi.h"
-#include "concretelang/Runtime/context.h"
-#include <stdio.h>
-
-LweKeyswitchKey_u64 *get_keyswitch_key(RuntimeContext *context) {
-  return context->ksk;
-}
-
-LweBootstrapKey_u64 *get_bootstrap_key(RuntimeContext *context) {
-  int err;
-  LweBootstrapKey_u64 *clone =
-    clone_lwe_bootstrap_key_u64(&err, context->bsk);
-  if (err != 0)
-    fprintf(stderr, "Runtime: cloning bootstrap key failed.\n");
-  return clone;
-}

compiler/lib/Runtime/context.cpp

@@ -0,0 +1,34 @@
+#include <assert.h>
+#include <concretelang/Runtime/context.h>
+#include <stdio.h>
+
+#ifdef CONCRETELANG_PARALLEL_EXECUTION_ENABLED
+#include <hpx/include/runtime.hpp>
+#endif
+
+LweKeyswitchKey_u64 *get_keyswitch_key(RuntimeContext *context) {
+  return context->ksk;
+}
+
+LweBootstrapKey_u64 *get_bootstrap_key(RuntimeContext *context) {
+  int err;
+#ifdef CONCRETELANG_PARALLEL_EXECUTION_ENABLED
+  std::string threadName = hpx::get_thread_name();
+  auto bskIt = context->bsk.find(threadName);
+  if (bskIt == context->bsk.end()) {
+    bskIt = context->bsk
+                .insert(std::pair<std::string, LweBootstrapKey_u64 *>(
+                    threadName,
+                    clone_lwe_bootstrap_key_u64(
+                        &err, context->bsk["_concretelang_base_context_bsk"])))
+                .first;
+    if (err != 0)
+      fprintf(stderr, "Runtime: cloning bootstrap key failed.\n");
+  }
+#else
+  std::string baseName = "_concretelang_base_context_bsk";
+  auto bskIt = context->bsk.find(baseName);
+#endif
+  assert(bskIt->second && "No bootstrap key available in context");
+  return bskIt->second;
+}

compiler/lib/Support/Jit.cpp

@@ -146,6 +146,10 @@ JITLambda::Argument::~Argument() {
   for (auto buffer : ciphertextBuffers) {
     free(buffer);
   }
+  for (const auto &[key, value] : context.bsk) {
+    if (key != "_concretelang_base_context_bsk")
+      free_lwe_bootstrap_key_u64(&err, value);
+  }
 }
 
 llvm::Expected<std::unique_ptr<JITLambda::Argument>>