wip

zama-ai · Oct 31, 2024 · 266da78 · 266da78
1 parent 7f47ee1
commit 266da78
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 122 deletions.
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,7 @@
 #ifndef CONCRETELANG_SUPPORT_V0Parameter_H_
 #define CONCRETELANG_SUPPORT_V0Parameter_H_
 
+#include <memory>
 #include <optional>
 #include <variant>
 
@@ -81,8 +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::optional<concrete_optimizer::restriction::RangeRestriction> DEFAULT_RANGE_RESTRICTION = std::nullopt;
-const std::optional<concrete_optimizer::restriction::KeysetRestriction> DEFAULT_KEYSET_RESTRICTION = std::nullopt;
+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 {
@@ -127,8 +128,8 @@ struct Config {
   bool cache_on_disk;
   uint32_t ciphertext_modulus_log;
   uint32_t fft_precision;
-  std::optional<concrete_optimizer::restriction::RangeRestriction> range_restriction;
-  std::optional<concrete_optimizer::restriction::KeysetRestriction> keyset_restriction;
+  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;
 };

diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp
@@ -340,6 +340,61 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
       .value("NATIVE", concrete_optimizer::Encoding::Native)
       .export_values();
 
+  pybind11::class_<concrete_optimizer::restriction::RangeRestriction>(
+      m, "RangeRestriction")
+      .def(pybind11::init(
+          []() { return concrete_optimizer::restriction::RangeRestriction(); }))
+      .def(
+          "add_available_glwe_log_polynomial_size",
+          [](concrete_optimizer::restriction::RangeRestriction &restriction,
+             uint64_t v) {
+            restriction.glwe_log_polynomial_sizes.push_back(v);
+          },
+          "Add an available glwe log poly size to the restriction")
+      .def(
+          "add_available_glwe_dimension",
+          [](concrete_optimizer::restriction::RangeRestriction &restriction,
+             uint64_t v) {
+            restriction.glwe_dimensions.push_back(v);
+          },
+          "Add an available glwe dimension to the restriction")
+      .def(
+          "add_available_internal_lwe_dimension",
+          [](concrete_optimizer::restriction::RangeRestriction &restriction,
+             uint64_t v) {
+            restriction.internal_lwe_dimensions.push_back(v);
+          },
+          "Add an available internal lwe dimension to the restriction")
+      .def(
+          "add_available_pbs_level_count",
+          [](concrete_optimizer::restriction::RangeRestriction &restriction,
+             uint64_t v) {
+            restriction.pbs_level_count.push_back(v);
+          },
+          "Add an available pbs level count to the restriction")
+      .def(
+          "add_available_pbs_base_log",
+          [](concrete_optimizer::restriction::RangeRestriction &restriction,
+             uint64_t v) {
+            restriction.pbs_base_log.push_back(v);
+          },
+          "Add an available pbs base log to the restriction")
+      .def(
+          "add_available_ks_level_count",
+          [](concrete_optimizer::restriction::RangeRestriction &restriction,
+             uint64_t v) {
+            restriction.ks_level_count.push_back(v);
+          },
+          "Add an available ks level count to the restriction")
+      .def(
+          "add_available_ks_base_log",
+          [](concrete_optimizer::restriction::RangeRestriction &restriction,
+             uint64_t v) {
+            restriction.ks_base_log.push_back(v);
+          },
+          "Add an available ks base log to the restriction")
+      .doc() = "Allow to restrict the optimizer parameter search space to a set of values.";
+
   pybind11::class_<CompilationOptions>(m, "CompilationOptions")
       .def(pybind11::init([](mlir::concretelang::Backend backend) {
              return CompilationOptions(backend);
@@ -444,117 +499,10 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
             options.optimizerConfig.security = security_level;
           },
           "Set security level.", arg("security_level"))
-      .def("set_glwe_pbs_restrictions",
-           [](CompilationOptions &options,
-              std::optional<uint64_t> log2_polynomial_size_min,
-              std::optional<uint64_t> log2_polynomial_size_max,
-              std::optional<uint64_t> glwe_dimension_min,
-              std::optional<uint64_t> glwe_dimension_max) {
-             if (log2_polynomial_size_min) {
-               options.optimizerConfig.parameter_restrictions.glwe_pbs
-                   .log2_polynomial_size_min =
-                   std::make_shared<uint64_t>(*log2_polynomial_size_min);
-             }
-             if (log2_polynomial_size_max) {
-               options.optimizerConfig.parameter_restrictions.glwe_pbs
-                   .log2_polynomial_size_max =
-                   std::make_shared<uint64_t>(*log2_polynomial_size_max);
-             }
-             if (glwe_dimension_min) {
-               options.optimizerConfig.parameter_restrictions.glwe_pbs
-                   .glwe_dimension_min =
-                   std::make_shared<uint64_t>(*glwe_dimension_min);
-             }
-             if (glwe_dimension_max) {
-               options.optimizerConfig.parameter_restrictions.glwe_pbs
-                   .glwe_dimension_max =
-                   std::make_shared<uint64_t>(*glwe_dimension_max);
-             }
-           })
-      .def("set_free_glwe_restrictions",
-           [](CompilationOptions &options,
-              std::optional<uint64_t> log2_polynomial_size_min,
-              std::optional<uint64_t> log2_polynomial_size_max,
-              std::optional<uint64_t> glwe_dimension_min,
-              std::optional<uint64_t> glwe_dimension_max) {
-             if (log2_polynomial_size_min) {
-               options.optimizerConfig.parameter_restrictions.free_glwe
-                   .log2_polynomial_size_min =
-                   std::make_shared<uint64_t>(*log2_polynomial_size_min);
-             }
-             if (log2_polynomial_size_max) {
-               options.optimizerConfig.parameter_restrictions.free_glwe
-                   .log2_polynomial_size_max =
-                   std::make_shared<uint64_t>(*log2_polynomial_size_max);
-             }
-             if (glwe_dimension_min) {
-               options.optimizerConfig.parameter_restrictions.free_glwe
-                   .glwe_dimension_min =
-                   std::make_shared<uint64_t>(*glwe_dimension_min);
-             }
-             if (glwe_dimension_max) {
-               options.optimizerConfig.parameter_restrictions.free_glwe
-                   .glwe_dimension_max =
-                   std::make_shared<uint64_t>(*glwe_dimension_max);
-             }
-           })
-      .def("set_br_decomposition_restrictions",
-           [](CompilationOptions &options,
-              std::optional<uint64_t> log2_base_min,
-              std::optional<uint64_t> log2_base_max,
-              std::optional<uint64_t> level_min,
-              std::optional<uint64_t> level_max) {
-             if (log2_base_min) {
-               options.optimizerConfig.parameter_restrictions.br_decomposition
-                   .log2_base_min = std::make_shared<uint64_t>(*log2_base_min);
-             }
-             if (log2_base_max) {
-               options.optimizerConfig.parameter_restrictions.br_decomposition
-                   .log2_base_max = std::make_shared<uint64_t>(*log2_base_max);
-             }
-             if (level_min) {
-               options.optimizerConfig.parameter_restrictions.br_decomposition
-                   .level_min = std::make_shared<uint64_t>(*level_min);
-             }
-             if (level_max) {
-               options.optimizerConfig.parameter_restrictions.br_decomposition
-                   .level_max = std::make_shared<uint64_t>(*level_max);
-             }
-           })
-      .def("set_ks_decomposition_restrictions",
+      .def("set_range_restriction",
            [](CompilationOptions &options,
-              std::optional<uint64_t> log2_base_min,
-              std::optional<uint64_t> log2_base_max,
-              std::optional<uint64_t> level_min,
-              std::optional<uint64_t> level_max) {
-             if (log2_base_min) {
-               options.optimizerConfig.parameter_restrictions.ks_decomposition
-                   .log2_base_min = std::make_shared<uint64_t>(*log2_base_min);
-             }
-             if (log2_base_max) {
-               options.optimizerConfig.parameter_restrictions.ks_decomposition
-                   .log2_base_max = std::make_shared<uint64_t>(*log2_base_max);
-             }
-             if (level_min) {
-               options.optimizerConfig.parameter_restrictions.ks_decomposition
-                   .level_min = std::make_shared<uint64_t>(*level_min);
-             }
-             if (level_max) {
-               options.optimizerConfig.parameter_restrictions.ks_decomposition
-                   .level_max = std::make_shared<uint64_t>(*level_max);
-             }
-           })
-      .def("set_free_lwe_restrictions",
-           [](CompilationOptions &options, std::optional<uint64_t> free_lwe_min,
-              std::optional<uint64_t> free_lwe_max) {
-             if (free_lwe_min) {
-               options.optimizerConfig.parameter_restrictions.free_lwe_min =
-                   std::make_shared<uint64_t>(*free_lwe_min);
-             }
-             if (free_lwe_max) {
-               options.optimizerConfig.parameter_restrictions.free_lwe_max =
-                   std::make_shared<uint64_t>(*free_lwe_max);
-             }
+              concrete_optimizer::restriction::RangeRestriction restriction) {
+                *options.optimizerConfig.range_restriction = restriction;
            })
       .def(
           "set_v0_parameter",

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;
 }
 

diff --git a/compilers/concrete-optimizer/concrete-optimizer-cpp/src/concrete-optimizer.rs b/compilers/concrete-optimizer/concrete-optimizer-cpp/src/concrete-optimizer.rs
@@ -662,14 +662,14 @@ impl Dag {
             ffi::MultiParamStrategy::ByPrecision | _ => PartitionCut::for_each_precision(&self.0),
         };
         let circuit_sol =
-            if !options.keyset_restriction.is_null() && !options.parameter_restriction.is_null() {
+            if !options.keyset_restriction.is_null() && !options.range_restriction.is_null() {
                 concrete_optimizer::optimization::dag::multi_parameters::optimize_generic::optimize(
                     &self.0,
                     config,
                     &search_space,
                     &(
                         (*options.keyset_restriction).clone(),
-                        (*options.parameter_restriction).clone(),
+                        (*options.range_restriction).clone(),
                     ),
                     encoding,
                     options.default_log_norm2_woppbs,
@@ -687,12 +687,12 @@ impl Dag {
                     &caches_from(options),
                     &Some(p_cut),
                 )
-            } else if !options.parameter_restriction.is_null() {
+            } else if !options.range_restriction.is_null() {
                 concrete_optimizer::optimization::dag::multi_parameters::optimize_generic::optimize(
                     &self.0,
                     config,
                     &search_space,
-                    &*options.parameter_restriction,
+                    &*options.range_restriction,
                     encoding,
                     options.default_log_norm2_woppbs,
                     &caches_from(options),
@@ -1209,7 +1209,7 @@ mod ffi {
         pub cache_on_disk: bool,
         pub ciphertext_modulus_log: u32,
         pub fft_precision: u32,
-        pub parameter_restriction: SharedPtr<RangeRestriction>, // SharedPtr used for Options since optionals are not available...
+        pub range_restriction: SharedPtr<RangeRestriction>, // SharedPtr used for Options since optionals are not available...
         pub keyset_restriction: SharedPtr<KeysetRestriction>, // SharedPtr used for Options since optionals are not available...
     }
 

diff --git a/compilers/concrete-optimizer/concrete-optimizer-cpp/src/cpp/concrete-optimizer.cpp b/compilers/concrete-optimizer/concrete-optimizer-cpp/src/cpp/concrete-optimizer.cpp
@@ -1175,7 +1175,7 @@ struct Options final {
   bool cache_on_disk;
   ::std::uint32_t ciphertext_modulus_log;
   ::std::uint32_t fft_precision;
-  ::std::shared_ptr<::concrete_optimizer::restriction::RangeRestriction> parameter_restriction;
+  ::std::shared_ptr<::concrete_optimizer::restriction::RangeRestriction> range_restriction;
   ::std::shared_ptr<::concrete_optimizer::restriction::KeysetRestriction> keyset_restriction;
 
   using IsRelocatable = ::std::true_type;

diff --git a/compilers/concrete-optimizer/concrete-optimizer-cpp/src/cpp/concrete-optimizer.hpp b/compilers/concrete-optimizer/concrete-optimizer-cpp/src/cpp/concrete-optimizer.hpp
@@ -1156,7 +1156,7 @@ struct Options final {
   bool cache_on_disk;
   ::std::uint32_t ciphertext_modulus_log;
   ::std::uint32_t fft_precision;
-  ::std::shared_ptr<::concrete_optimizer::restriction::RangeRestriction> parameter_restriction;
+  ::std::shared_ptr<::concrete_optimizer::restriction::RangeRestriction> range_restriction;
   ::std::shared_ptr<::concrete_optimizer::restriction::KeysetRestriction> keyset_restriction;
 
   using IsRelocatable = ::std::true_type;