Merge pull request #1120 from zama-ai/custom_parameter_restrictions

feat(optimizer): add ability to configure custom parameter restrictions
zama-ai · Oct 28, 2024 · b64fa7c · b64fa7c
2 parents bd36a47 + 019daf8
commit b64fa7c
Show file tree

Hide file tree

Showing 15 changed files with 586 additions and 88 deletions.
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Support/V0Parameters.h b/compilers/concrete-compiler/compiler/include/concretelang/Support/V0Parameters.h
@@ -124,6 +124,7 @@ struct Config {
   bool cache_on_disk;
   uint32_t ciphertext_modulus_log;
   uint32_t fft_precision;
+  concrete_optimizer::ParameterRestrictions parameter_restrictions;
   std::vector<CompositionRule> composition_rules;
   bool composable;
 };
@@ -141,6 +142,7 @@ const Config DEFAULT_CONFIG = {UNSPECIFIED_P_ERROR,
                                DEFAULT_CACHE_ON_DISK,
                                DEFAULT_CIPHERTEXT_MODULUS_LOG,
                                DEFAULT_FFT_PRECISION,
+                               concrete_optimizer::ParameterRestrictions{},
                                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
@@ -444,6 +444,118 @@ 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",
+           [](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);
+             }
+           })
       .def(
           "set_v0_parameter",
           [](CompilationOptions &options, size_t glweDimension,

diff --git a/compilers/concrete-compiler/compiler/lib/Support/V0Parameters.cpp b/compilers/concrete-compiler/compiler/lib/Support/V0Parameters.cpp
@@ -35,7 +35,8 @@ concrete_optimizer::Options options_from_config(optimizer::Config config) {
       /* .encoding = */ config.encoding,
       /* .cache_on_disk = */ config.cache_on_disk,
       /* .ciphertext_modulus_log = */ config.ciphertext_modulus_log,
-      /* .fft_precision = */ config.fft_precision};
+      /* .fft_precision = */ config.fft_precision,
+      /* .parameter_restrictions = */ config.parameter_restrictions};
   return options;
 }
 

diff --git a/compilers/concrete-optimizer/charts/src/bin/norm2_complexity.rs b/compilers/concrete-optimizer/charts/src/bin/norm2_complexity.rs
@@ -1,21 +1,25 @@
 use charts::{draw, Serie};
 use concrete_optimizer::computing_cost::cpu::CpuComplexity;
 use concrete_optimizer::config;
-use concrete_optimizer::global_parameters::DEFAUT_DOMAINS;
+use concrete_optimizer::global_parameters::DEFAULT_DOMAINS;
 use concrete_optimizer::optimization::atomic_pattern::{self as optimize_atomic_pattern};
 use concrete_optimizer::optimization::config::{Config, SearchSpace};
 use concrete_optimizer::optimization::decomposition;
 use concrete_optimizer::optimization::wop_atomic_pattern::optimize as optimize_wop_atomic_pattern;
 
 pub const _4_SIGMA: f64 = 1.0 - 0.999_936_657_516;
-const MIN_LOG_POLY_SIZE: u64 = DEFAUT_DOMAINS
-    .glwe_pbs_constrained
+const MIN_LOG_POLY_SIZE: u64 = DEFAULT_DOMAINS
+    .glwe_pbs_constrained_cpu
     .log2_polynomial_size
     .start;
-const MAX_LOG_POLY_SIZE: u64 = DEFAUT_DOMAINS.glwe_pbs_constrained.log2_polynomial_size.end - 1;
-pub const MAX_GLWE_DIM: u64 = DEFAUT_DOMAINS.glwe_pbs_constrained.glwe_dimension.end - 1;
-pub const MIN_LWE_DIM: u64 = DEFAUT_DOMAINS.free_glwe.glwe_dimension.start;
-pub const MAX_LWE_DIM: u64 = DEFAUT_DOMAINS.free_glwe.glwe_dimension.end - 1;
+const MAX_LOG_POLY_SIZE: u64 = DEFAULT_DOMAINS
+    .glwe_pbs_constrained_cpu
+    .log2_polynomial_size
+    .end
+    - 1;
+pub const MAX_GLWE_DIM: u64 = DEFAULT_DOMAINS.glwe_pbs_constrained_cpu.glwe_dimension.end - 1;
+pub const MIN_LWE_DIM: u64 = DEFAULT_DOMAINS.free_glwe.glwe_dimension.start;
+pub const MAX_LWE_DIM: u64 = DEFAULT_DOMAINS.free_glwe.glwe_dimension.end - 1;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let sum_size = 4096;
@@ -31,7 +35,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         glwe_log_polynomial_sizes,
         glwe_dimensions,
         internal_lwe_dimensions,
-        levelled_only_lwe_dimensions: DEFAUT_DOMAINS.free_lwe,
+        levelled_only_lwe_dimensions: DEFAULT_DOMAINS.free_lwe,
     };
 
     let precision = 8;

diff --git a/compilers/concrete-optimizer/charts/src/bin/precision_complexity.rs b/compilers/concrete-optimizer/charts/src/bin/precision_complexity.rs
@@ -1,21 +1,25 @@
 use charts::{draw, Serie};
 use concrete_optimizer::computing_cost::cpu::CpuComplexity;
 use concrete_optimizer::config;
-use concrete_optimizer::global_parameters::DEFAUT_DOMAINS;
+use concrete_optimizer::global_parameters::DEFAULT_DOMAINS;
 use concrete_optimizer::optimization::atomic_pattern::{self as optimize_atomic_pattern};
 use concrete_optimizer::optimization::config::{Config, SearchSpace};
 use concrete_optimizer::optimization::decomposition;
 use concrete_optimizer::optimization::wop_atomic_pattern::optimize as optimize_wop_atomic_pattern;
 
 pub const _4_SIGMA: f64 = 1.0 - 0.999_936_657_516;
-const MIN_LOG_POLY_SIZE: u64 = DEFAUT_DOMAINS
-    .glwe_pbs_constrained
+const MIN_LOG_POLY_SIZE: u64 = DEFAULT_DOMAINS
+    .glwe_pbs_constrained_cpu
     .log2_polynomial_size
     .start;
-const MAX_LOG_POLY_SIZE: u64 = DEFAUT_DOMAINS.glwe_pbs_constrained.log2_polynomial_size.end - 1;
-pub const MAX_GLWE_DIM: u64 = DEFAUT_DOMAINS.glwe_pbs_constrained.glwe_dimension.end - 1;
-pub const MIN_LWE_DIM: u64 = DEFAUT_DOMAINS.free_glwe.glwe_dimension.start;
-pub const MAX_LWE_DIM: u64 = DEFAUT_DOMAINS.free_glwe.glwe_dimension.end - 1;
+const MAX_LOG_POLY_SIZE: u64 = DEFAULT_DOMAINS
+    .glwe_pbs_constrained_cpu
+    .log2_polynomial_size
+    .end
+    - 1;
+pub const MAX_GLWE_DIM: u64 = DEFAULT_DOMAINS.glwe_pbs_constrained_cpu.glwe_dimension.end - 1;
+pub const MIN_LWE_DIM: u64 = DEFAULT_DOMAINS.free_glwe.glwe_dimension.start;
+pub const MAX_LWE_DIM: u64 = DEFAULT_DOMAINS.free_glwe.glwe_dimension.end - 1;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let processing_unit = config::ProcessingUnit::Cpu;
@@ -37,7 +41,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         glwe_log_polynomial_sizes,
         glwe_dimensions,
         internal_lwe_dimensions,
-        levelled_only_lwe_dimensions: DEFAUT_DOMAINS.free_lwe,
+        levelled_only_lwe_dimensions: DEFAULT_DOMAINS.free_lwe,
     };
 
     let config = Config {

diff --git a/compilers/concrete-optimizer/concrete-optimizer-cpp/Makefile b/compilers/concrete-optimizer/concrete-optimizer-cpp/Makefile
@@ -37,7 +37,7 @@ $(INTERFACE_LIB): $(INTERFACE_LIB_ORIG)
 TESTS_SOURCES = tests/src/main.cpp
 TEST_DEP_LIBS = -l pthread -ldl
 tests/tests_exe: $(INTERFACE_LIB) $(INTERFACE_HEADER) $(INTERFACE_CPP) $(TESTS_SOURCES)
-	g++ -Wall -Werror -Wextra -o $@ $(TESTS_SOURCES) $(INTERFACE_CPP) $(INTERFACE_LIB) -I $(shell dirname $(INTERFACE_HEADER)) $(TEST_DEP_LIBS)
+	g++ -Wall -Werror -Wextra -std=c++17 -o $@ $(TESTS_SOURCES) $(INTERFACE_CPP) $(INTERFACE_LIB) -I $(shell dirname $(INTERFACE_HEADER)) $(TEST_DEP_LIBS)
 	chmod +x $@
 
 test: tests/tests_exe