feat(compiler): search space restriction

zama-ai · Nov 12, 2024 · 0fccfbf · 0fccfbf
1 parent 81af82f
commit 0fccfbf
Show file tree

Hide file tree

Showing 20 changed files with 1,545 additions and 593 deletions.
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Common/Protocol.h b/compilers/concrete-compiler/compiler/include/concretelang/Common/Protocol.h
@@ -114,6 +114,18 @@ template <typename MessageType> struct Message {
     return *this;
   }
 
+  bool operator==(Message const& other) const {
+      capnp::AnyStruct::Reader left = this->asReader();
+      capnp::AnyStruct::Reader right = other.asReader();
+      return left == right;
+  }
+
+  bool operator!=(Message const& other) const {
+      capnp::AnyStruct::Reader left = this->asReader();
+      capnp::AnyStruct::Reader right = other.asReader();
+      return left != right;
+  }
+
   Message(Message &&input) : message(nullptr) {
     regionBuilder = input.regionBuilder;
     message = input.message;

diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Support/V0Parameters.h b/compilers/concrete-compiler/compiler/include/concretelang/Support/V0Parameters.h
@@ -6,6 +6,8 @@
 #ifndef CONCRETELANG_SUPPORT_V0Parameter_H_
 #define CONCRETELANG_SUPPORT_V0Parameter_H_
 
+#include <memory>
+#include <optional>
 #include <variant>
 
 #include "llvm/ADT/Optional.h"
@@ -80,6 +82,8 @@ const concrete_optimizer::Encoding DEFAULT_ENCODING =
 const bool DEFAULT_CACHE_ON_DISK = true;
 const uint32_t DEFAULT_CIPHERTEXT_MODULUS_LOG = 64;
 const uint32_t DEFAULT_FFT_PRECISION = 53;
+const std::shared_ptr<concrete_optimizer::restriction::RangeRestriction> DEFAULT_RANGE_RESTRICTION = {};
+const std::shared_ptr<concrete_optimizer::restriction::KeysetRestriction> DEFAULT_KEYSET_RESTRICTION = {};
 
 /// The strategy of the crypto optimization
 enum Strategy {
@@ -124,7 +128,8 @@ struct Config {
   bool cache_on_disk;
   uint32_t ciphertext_modulus_log;
   uint32_t fft_precision;
-  concrete_optimizer::ParameterRestrictions parameter_restrictions;
+  std::shared_ptr<concrete_optimizer::restriction::RangeRestriction> range_restriction;
+  std::shared_ptr<concrete_optimizer::restriction::KeysetRestriction> keyset_restriction;
   std::vector<CompositionRule> composition_rules;
   bool composable;
 };
@@ -142,7 +147,8 @@ const Config DEFAULT_CONFIG = {UNSPECIFIED_P_ERROR,
                                DEFAULT_CACHE_ON_DISK,
                                DEFAULT_CIPHERTEXT_MODULUS_LOG,
                                DEFAULT_FFT_PRECISION,
-                               concrete_optimizer::ParameterRestrictions{},
+                               DEFAULT_RANGE_RESTRICTION,
+                               DEFAULT_KEYSET_RESTRICTION,
                                DEFAULT_COMPOSITION_RULES,
                                DEFAULT_COMPOSABLE};
 

diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp
diff --git a/.../concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_feedback.py b/.../concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_feedback.py
@@ -100,13 +100,13 @@ def count_per_parameter(
                     continue
 
                 if key_type == KeyType.SECRET:
-                    parameter = program_info.secret_keys()[key_index]
+                    parameter = program_info.get_keyset_info().secret_keys()[key_index]
                 elif key_type == KeyType.BOOTSTRAP:
-                    parameter = program_info.bootstrap_keys()[key_index]
+                    parameter = program_info.get_keyset_info().bootstrap_keys()[key_index]
                 elif key_type == KeyType.KEY_SWITCH:
-                    parameter = program_info.keyswitch_keys()[key_index]
+                    parameter = program_info.get_keyset_info().keyswitch_keys()[key_index]
                 elif key_type == KeyType.PACKING_KEY_SWITCH:
-                    parameter = program_info.packing_keyswitch_keys()[key_index]
+                    parameter = program_info.get_keyset_info().packing_keyswitch_keys()[key_index]
                 else:
                     assert False
                 if parameter not in result:

diff --git a/compilers/concrete-compiler/compiler/lib/Support/V0Parameters.cpp b/compilers/concrete-compiler/compiler/lib/Support/V0Parameters.cpp
@@ -36,7 +36,15 @@ concrete_optimizer::Options options_from_config(optimizer::Config config) {
       /* .cache_on_disk = */ config.cache_on_disk,
       /* .ciphertext_modulus_log = */ config.ciphertext_modulus_log,
       /* .fft_precision = */ config.fft_precision,
-      /* .parameter_restrictions = */ config.parameter_restrictions};
+      /* .range_restriction = */ std::shared_ptr<concrete_optimizer::restriction::RangeRestriction>(),
+      /* .keyset_restriction = */ std::shared_ptr<concrete_optimizer::restriction::KeysetRestriction>(),
+  };
+  if (config.range_restriction){
+      options.range_restriction = config.range_restriction;
+  }
+  if (config.keyset_restriction){
+      options.keyset_restriction = config.keyset_restriction;
+  }
   return options;
 }