feat(optimizer): fetch caches earlier

concrete-optimizer/src/optimization/atomic_pattern.rs

@@ -1,14 +1,14 @@
 use concrete_cpu_noise_model::gaussian_noise::noise::modulus_switching::estimate_modulus_switching_noise_with_binary_key;
 
 use super::config::{Config, SearchSpace};
+use super::decomposition::cmux::CmuxComplexityNoise;
+use super::decomposition::keyswitch::KsComplexityNoise;
 use super::wop_atomic_pattern::optimize::find_p_error;
 use crate::noise_estimator::error;
 use crate::parameters::{BrDecompositionParameters, GlweParameters, KsDecompositionParameters};
 use crate::utils::square;
 
-use super::decomposition::{
-    circuit_bootstrap, cmux, keyswitch, pp_switch, DecompCaches, PersistDecompCaches,
-};
+use super::decomposition::{circuit_bootstrap, cmux, keyswitch, pp_switch, PersistDecompCaches};
 
 // Ref time for v0 table 1 thread: 950ms
 const CUTS: bool = true; // 80ms
@@ -55,7 +55,8 @@ fn update_state_with_best_decompositions(
     consts: &OptimizationDecompositionsConsts,
     internal_dim: u64,
     glwe_params: GlweParameters,
-    caches: &mut DecompCaches,
+    cmux_quantities: &[CmuxComplexityNoise],
+    ks_quantities: &[KsComplexityNoise],
 ) {
     let input_lwe_dimension = glwe_params.sample_extract_lwe_dimension();
     let noise_modulus_switching = estimate_modulus_switching_noise_with_binary_key(
@@ -72,9 +73,6 @@ fn update_state_with_best_decompositions(
     let mut best_variance = state.best_solution.map_or(f64::INFINITY, |s| s.noise_max);
 
     let complexity_multisum = (consts.sum_size * input_lwe_dimension) as f64;
-    let cmux_quantities = caches.cmux.pareto_quantities(glwe_params);
-
-    let ks_quantities = caches.keyswitch.pareto_quantities(internal_dim);
 
     let square_noise_factor = square(consts.noise_factor);
     for cmux_quantity in cmux_quantities {
@@ -204,14 +202,20 @@ pub fn optimize_one(
                 glwe_dimension: glwe_dim,
             };
 
+            let cmux_quantities = caches.cmux.pareto_quantities(glwe_params);
+
             for &internal_dim in &search_space.internal_lwe_dimensions {
                 assert!(256 < internal_dim);
+
+                let ks_quantities = caches.keyswitch.pareto_quantities(internal_dim);
+
                 update_state_with_best_decompositions(
                     &mut state,
                     &consts,
                     internal_dim,
                     glwe_params,
-                    &mut caches,
+                    cmux_quantities,
+                    ks_quantities,
                 );
             }
         }

concrete-optimizer/src/optimization/dag/solo_key/optimize.rs

@@ -8,7 +8,9 @@ use crate::optimization::atomic_pattern::{
     OptimizationDecompositionsConsts, OptimizationState, Solution,
 };
 use crate::optimization::config::{Config, NoiseBoundConfig, SearchSpace};
-use crate::optimization::decomposition::{DecompCaches, PersistDecompCaches};
+use crate::optimization::decomposition::cmux::CmuxComplexityNoise;
+use crate::optimization::decomposition::keyswitch::KsComplexityNoise;
+use crate::optimization::decomposition::PersistDecompCaches;
 use crate::parameters::GlweParameters;
 
 #[allow(clippy::too_many_lines)]
@@ -20,7 +22,8 @@ fn update_best_solution_with_best_decompositions(
     glwe_params: GlweParameters,
     input_noise_out: f64,
     noise_modulus_switching: f64,
-    caches: &mut DecompCaches,
+    cmux_pareto: &[CmuxComplexityNoise],
+    ks_pareto: &[KsComplexityNoise],
 ) {
     assert!(dag.nb_luts > 0);
     let input_lwe_dimension = glwe_params.sample_extract_lwe_dimension();
@@ -29,10 +32,6 @@ fn update_best_solution_with_best_decompositions(
     let mut best_variance = state.best_solution.map_or(f64::INFINITY, |s| s.noise_max);
     let mut best_p_error = state.best_solution.map_or(f64::INFINITY, |s| s.p_error);
 
-    let cmux_pareto = caches.cmux.pareto_quantities(glwe_params);
-
-    let ks_pareto = caches.keyswitch.pareto_quantities(internal_dim);
-
     // by constructon br_pareto and ks_pareto are non-empty
     let mut best_cmux = cmux_pareto[0];
     let mut best_ks = ks_pareto[0];
@@ -282,7 +281,12 @@ pub fn optimize(
                 glwe_dimension: glwe_dim,
             };
             let input_noise_out = minimal_variance(&config, glwe_params);
+
+            let cmux_pareto = caches.cmux.pareto_quantities(glwe_params);
+
             for &internal_dim in &search_space.internal_lwe_dimensions {
+                let ks_pareto = caches.keyswitch.pareto_quantities(internal_dim);
+
                 let noise_modulus_switching =
                     noise_modulus_switching(glwe_log_poly_size, internal_dim);
                 if not_feasible(input_noise_out, noise_modulus_switching) {
@@ -297,7 +301,8 @@ pub fn optimize(
                     glwe_params,
                     input_noise_out,
                     noise_modulus_switching,
-                    &mut caches,
+                    cmux_pareto,
+                    ks_pareto,
                 );
                 if dag.nb_luts == 0 && state.best_solution.is_some() {
                     return state;

concrete-optimizer/src/optimization/wop_atomic_pattern/optimize.rs

@@ -14,7 +14,8 @@ use crate::optimization::config::{Config, SearchSpace};
 use crate::optimization::decomposition::circuit_bootstrap::CbComplexityNoise;
 use crate::optimization::decomposition::cmux::CmuxComplexityNoise;
 use crate::optimization::decomposition::keyswitch::KsComplexityNoise;
-use crate::optimization::decomposition::{DecompCaches, PersistDecompCaches};
+use crate::optimization::decomposition::pp_switch::PpSwitchComplexityNoise;
+use crate::optimization::decomposition::PersistDecompCaches;
 use crate::parameters::{BrDecompositionParameters, GlweParameters};
 use crate::utils::square;
 
@@ -193,7 +194,10 @@ fn update_state_with_best_decompositions(
     internal_dim: u64,
     n_functions: u64,
     partitionning: &[u64],
-    caches: &mut DecompCaches,
+    pareto_cmux: &[CmuxComplexityNoise],
+    pareto_keyswitch: &[KsComplexityNoise],
+    pp_switch: &[PpSwitchComplexityNoise],
+    pareto_cb: &[CbComplexityNoise],
 ) {
     let ciphertext_modulus_log = consts.config.ciphertext_modulus_log;
     let precisions_sum = partitionning.iter().copied().sum();
@@ -224,14 +228,6 @@ fn update_state_with_best_decompositions(
         .as_ref()
         .map_or(f64::INFINITY, |s| s.noise_max);
 
-    let pareto_cmux = caches.cmux.pareto_quantities(glwe_params);
-
-    let pareto_keyswitch = caches.keyswitch.pareto_quantities(internal_dim);
-
-    let pp_switch = caches.pp_switch.pareto_quantities(glwe_params);
-
-    let pareto_cb = caches.cb_pbs.pareto_quantities(glwe_params);
-
     let lower_bound_variance_br = pareto_cmux.last().unwrap().noise_br(internal_dim);
     let lower_bound_variance_ks = pareto_keyswitch.last().unwrap().noise(input_lwe_dimension);
     let lower_bound_variance_private_packing = pp_switch.last().unwrap().noise;
@@ -483,7 +479,15 @@ fn optimize_raw(
                 glwe_dimension: glwe_dim,
             };
 
+            let pareto_cmux = caches.cmux.pareto_quantities(glwe_params);
+
+            let pareto_pp_switch = caches.pp_switch.pareto_quantities(glwe_params);
+
+            let pareto_cb = caches.cb_pbs.pareto_quantities(glwe_params);
+
             for &internal_dim in &search_space.internal_lwe_dimensions {
+                let pareto_keyswitch = caches.keyswitch.pareto_quantities(internal_dim);
+
                 update_state_with_best_decompositions(
                     &mut state,
                     &consts,
@@ -491,7 +495,10 @@ fn optimize_raw(
                     internal_dim,
                     n_functions,
                     partitionning,
-                    &mut caches,
+                    pareto_cmux,
+                    pareto_keyswitch,
+                    pareto_pp_switch,
+                    pareto_cb,
                 );
             }
         }