Merge pull request #1125 from zama-ai/alex/optimizer_tfhers

Search space restriction
zama-ai · Nov 15, 2024 · b796dc8 · b796dc8
2 parents c5d1bf7 + f70a18b
commit b796dc8
Show file tree

Hide file tree

Showing 20 changed files with 1,577 additions and 597 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,10 @@ 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 +130,10 @@ 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 +151,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,19 @@ 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:
@@ -201,13 +207,23 @@ def count_per_tag_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
 

diff --git a/compilers/concrete-compiler/compiler/lib/Support/V0Parameters.cpp b/compilers/concrete-compiler/compiler/lib/Support/V0Parameters.cpp
@@ -36,7 +36,17 @@ 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;
 }