fix(compiler): Share the decompress state instead of try to avoid copies

This partially reverts commit 081e8b7.
zama-ai · May 14, 2024 · efe223e · efe223e
1 parent c740503
commit efe223e
Show file tree

Hide file tree

Showing 7 changed files with 22 additions and 20 deletions.
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Common/Keys.h b/compilers/concrete-compiler/compiler/include/concretelang/Common/Keys.h
@@ -89,7 +89,7 @@ class LweBootstrapKey {
                   Message<concreteprotocol::LweBootstrapKeyInfo> info)
       : seededBuffer(std::make_shared<std::vector<uint64_t>>()), buffer(buffer),
         info(info), decompress_mutext(std::make_shared<std::mutex>()),
-        decompressed(false){};
+        decompressed(std::make_shared<bool>(false)){};
 
   /// @brief Initialize the key from the protocol message.
   static LweBootstrapKey
@@ -111,7 +111,7 @@ class LweBootstrapKey {
       : seededBuffer(std::make_shared<std::vector<uint64_t>>()),
         buffer(std::make_shared<std::vector<uint64_t>>()), info(info),
         decompress_mutext(std::make_shared<std::mutex>()),
-        decompressed(false){};
+        decompressed(std::make_shared<bool>(false)){};
   LweBootstrapKey() = delete;
 
   /// @brief  The buffer of the seeded key if needed.
@@ -127,7 +127,7 @@ class LweBootstrapKey {
   std::shared_ptr<std::mutex> decompress_mutext;
 
   /// @brief A boolean that indicates if the decompression is done or not
-  bool decompressed;
+  std::shared_ptr<bool> decompressed;
 };
 
 class LweKeyswitchKey {
@@ -141,7 +141,7 @@ class LweKeyswitchKey {
                   Message<concreteprotocol::LweKeyswitchKeyInfo> info)
       : seededBuffer(std::make_shared<std::vector<uint64_t>>()), buffer(buffer),
         info(info), decompress_mutext(std::make_shared<std::mutex>()),
-        decompressed(false){};
+        decompressed(std::make_shared<bool>(false)){};
 
   /// @brief Initialize the key from the protocol message.
   static LweKeyswitchKey
@@ -163,7 +163,7 @@ class LweKeyswitchKey {
       : seededBuffer(std::make_shared<std::vector<uint64_t>>()),
         buffer(std::make_shared<std::vector<uint64_t>>()), info(info),
         decompress_mutext(std::make_shared<std::mutex>()),
-        decompressed(false){};
+        decompressed(std::make_shared<bool>(false)){};
 
   /// @brief  The buffer of the seeded key if needed.
   std::shared_ptr<std::vector<uint64_t>> seededBuffer;
@@ -178,7 +178,7 @@ class LweKeyswitchKey {
   std::shared_ptr<std::mutex> decompress_mutext;
 
   /// @brief A boolean that indicates if the decompression is done or not
-  bool decompressed;
+  std::shared_ptr<bool> decompressed;
 };
 
 class PackingKeyswitchKey {

diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Runtime/context.h b/compilers/concrete-compiler/compiler/include/concretelang/Runtime/context.h
@@ -42,7 +42,7 @@ typedef struct FFT {
 typedef struct RuntimeContext {
 
   RuntimeContext() = delete;
-  RuntimeContext(ServerKeyset &serverKeyset);
+  RuntimeContext(ServerKeyset serverKeyset);
   virtual ~RuntimeContext() {
 #ifdef CONCRETELANG_CUDA_SUPPORT
     for (int i = 0; i < num_devices; ++i) {
@@ -71,8 +71,10 @@ typedef struct RuntimeContext {
 
   virtual const struct Fft *fft(size_t keyId) { return ffts[keyId].fft; }
 
+  const ServerKeyset getKeys() const { return serverKeyset; }
+
 protected:
-  ServerKeyset &serverKeyset;
+  ServerKeyset serverKeyset;
   std::vector<std::shared_ptr<std::vector<std::complex<double>>>>
       fourier_bootstrap_keys;
   std::vector<FFT> ffts;

diff --git a/compilers/concrete-compiler/compiler/include/concretelang/ServerLib/ServerLib.h b/compilers/concrete-compiler/compiler/include/concretelang/ServerLib/ServerLib.h
@@ -47,7 +47,7 @@ class ServerCircuit {
 
 public:
   /// Call the circuit with public arguments.
-  Result<std::vector<TransportValue>> call(ServerKeyset &serverKeyset,
+  Result<std::vector<TransportValue>> call(const ServerKeyset &serverKeyset,
                                            std::vector<TransportValue> &args);
 
   /// Simulate the circuit with public arguments.
@@ -65,7 +65,7 @@ class ServerCircuit {
                     std::shared_ptr<DynamicModule> dynamicModule,
                     bool useSimulation);
 
-  void invoke(ServerKeyset &serverKeyset);
+  void invoke(const ServerKeyset &serverKeyset);
 
   Message<concreteprotocol::CircuitInfo> circuitInfo;
   bool useSimulation;

diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp
@@ -1226,7 +1226,7 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
               ::concretelang::clientlib::PublicArguments &publicArguments,
               ::concretelang::clientlib::EvaluationKeys &evaluationKeys) {
              SignalGuard signalGuard;
-             auto &keyset = evaluationKeys.keyset;
+             auto keyset = evaluationKeys.keyset;
              auto values = publicArguments.values;
              GET_OR_THROW_RESULT(auto output, circuit.call(keyset, values));
              ::concretelang::clientlib::PublicResult res{output};

diff --git a/compilers/concrete-compiler/compiler/lib/Common/Keys.cpp b/compilers/concrete-compiler/compiler/lib/Common/Keys.cpp
@@ -189,10 +189,10 @@ void LweBootstrapKey::decompress() {
   case concreteprotocol::Compression::NONE:
     return;
   case concreteprotocol::Compression::SEED: {
-    if (decompressed)
+    if (*decompressed)
       return;
     const std::lock_guard<std::mutex> guard(*decompress_mutext);
-    if (decompressed)
+    if (*decompressed)
       return;
     auto params = info.asReader().getParams();
     buffer->resize(concrete_cpu_bootstrap_key_size_u64(
@@ -204,7 +204,7 @@ void LweBootstrapKey::decompress() {
         buffer->data(), seededBuffer->data() + 2, params.getInputLweDimension(),
         params.getPolynomialSize(), params.getGlweDimension(),
         params.getLevelCount(), params.getBaseLog(), seed, Parallelism::Rayon);
-    decompressed = true;
+    *decompressed = true;
     return;
   }
   default:
@@ -309,10 +309,10 @@ void LweKeyswitchKey::decompress() {
   case concreteprotocol::Compression::NONE:
     return;
   case concreteprotocol::Compression::SEED: {
-    if (decompressed)
+    if (*decompressed)
       return;
     const std::lock_guard<std::mutex> guard(*decompress_mutext);
-    if (decompressed)
+    if (*decompressed)
       return;
     auto params = info.asReader().getParams();
     buffer->resize(concrete_cpu_keyswitch_key_size_u64(
@@ -324,7 +324,7 @@ void LweKeyswitchKey::decompress() {
         buffer->data(), seededBuffer->data() + 2, params.getInputLweDimension(),
         params.getOutputLweDimension(), params.getLevelCount(),
         params.getBaseLog(), seed, Parallelism::Rayon);
-    decompressed = true;
+    *decompressed = true;
     return;
   }
   default:

diff --git a/compilers/concrete-compiler/compiler/lib/Runtime/context.cpp b/compilers/concrete-compiler/compiler/lib/Runtime/context.cpp
@@ -29,7 +29,7 @@ FFT::~FFT() {
   }
 }
 
-RuntimeContext::RuntimeContext(ServerKeyset &serverKeyset)
+RuntimeContext::RuntimeContext(ServerKeyset serverKeyset)
     : serverKeyset(serverKeyset) {
 
   // Initialize for each bootstrap key the fourier one

diff --git a/compilers/concrete-compiler/compiler/lib/ServerLib/ServerLib.cpp b/compilers/concrete-compiler/compiler/lib/ServerLib/ServerLib.cpp
@@ -428,7 +428,7 @@ bool getGateIsSigned(const Message<concreteprotocol::GateInfo> &gateInfo) {
 }
 
 Result<std::vector<TransportValue>>
-ServerCircuit::call(ServerKeyset &serverKeyset,
+ServerCircuit::call(const ServerKeyset &serverKeyset,
                     std::vector<TransportValue> &args) {
   if (args.size() != argsBuffer.size()) {
     return StringError("Called circuit with wrong number of arguments");
@@ -542,7 +542,7 @@ Result<ServerCircuit> ServerCircuit::fromDynamicModule(
   return output;
 }
 
-void ServerCircuit::invoke(ServerKeyset &serverKeyset) {
+void ServerCircuit::invoke(const ServerKeyset &serverKeyset) {
 
   // We create a runtime context from the keyset, and place a pointer to it in
   // the structure.