fix(compiler): Decompression of evaluation keys was done several times

zama-ai · Apr 29, 2024 · fd2ea3f · fd2ea3f
1 parent 3d338af
commit fd2ea3f
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 10 deletions.
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) {
@@ -69,10 +69,8 @@ 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(const ServerKeyset &serverKeyset,
+  Result<std::vector<TransportValue>> call(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(const ServerKeyset &serverKeyset);
+  void invoke(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/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(const ServerKeyset &serverKeyset,
+ServerCircuit::call(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(const ServerKeyset &serverKeyset) {
+void ServerCircuit::invoke(ServerKeyset &serverKeyset) {
 
   // We create a runtime context from the keyset, and place a pointer to it in
   // the structure.