From 311a24a0e21d09f8faebc104255b0b3730d3ea68 Mon Sep 17 00:00:00 2001 From: rudy Date: Tue, 19 Sep 2023 18:42:00 +0200 Subject: [PATCH] feat(optimizer): multi-parameters, direct variance and cost operation bound --- .../dag/multi_parameters/complexity.rs | 39 +++++++++ .../dag/multi_parameters/feasible.rs | 85 +++++++++++++++++++ .../dag/multi_parameters/operations_value.rs | 7 ++ .../dag/multi_parameters/optimize.rs | 42 ++++++--- 4 files changed, 161 insertions(+), 12 deletions(-) diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/complexity.rs b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/complexity.rs index 68490701c8..ccefa6adc4 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/complexity.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/complexity.rs @@ -74,4 +74,43 @@ impl Complexity { pub fn complexity(&self, costs: &OperationsValue) -> f64 { f64_dot(&self.counts, costs) } + + #[inline(never)] + pub fn ks_max_cost( + &self, + complexity_cut: f64, + costs: &OperationsValue, + src_partition: usize, + dst_partition: usize, + ) -> f64 { + let ks_index = costs.index.keyswitch_to_small(src_partition, dst_partition); + let actual_ks_cost = costs.values[ks_index]; + let ks_coeff = self.counts[self + .counts + .index + .keyswitch_to_small(src_partition, dst_partition)]; + let actual_complexity = self.complexity(costs) - ks_coeff * actual_ks_cost; + + (complexity_cut - actual_complexity) / ks_coeff + } + + #[inline(never)] + pub fn fks_max_cost( + &self, + complexity_cut: f64, + costs: &OperationsValue, + src_partition: usize, + dst_partition: usize, + ) -> f64 { + let fks_index = costs.index.keyswitch_to_big(src_partition, dst_partition); + let actual_fks_cost = costs.values[fks_index]; + let fks_coeff = self.counts[self + .counts + .index + .keyswitch_to_big(src_partition, dst_partition)]; + let actual_complexity = self.complexity(costs) - fks_coeff * actual_fks_cost; + + (complexity_cut - actual_complexity) / fks_coeff + } + } diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/feasible.rs b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/feasible.rs index fe67a1426c..f4786e0e3e 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/feasible.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/feasible.rs @@ -25,6 +25,91 @@ impl Feasible { } } + #[inline(never)] + pub fn pbs_max_feasible_variance( + &self, + operations_variance: &OperationsValue, + partition: usize, + ) -> f64 { + let pbs_index = operations_variance.index.pbs(partition); + let actual_pbs_variance = operations_variance.values[pbs_index]; + + let mut smallest_pbs_max_variance = std::f64::MAX; + + for constraint in &self.undominated_constraints { + let pbs_coeff = constraint.variance.coeff_pbs(partition); + if pbs_coeff == 0.0 { + continue; + } + let actual_variance = f64_dot(operations_variance, &constraint.variance.coeffs) + - pbs_coeff * actual_pbs_variance; + let pbs_max_variance = (constraint.safe_variance_bound - actual_variance) / pbs_coeff; + smallest_pbs_max_variance = smallest_pbs_max_variance.min(pbs_max_variance); + } + smallest_pbs_max_variance + } + + #[inline(never)] + pub fn ks_max_feasible_variance( + &self, + operations_variance: &OperationsValue, + src_partition: usize, + dst_partition: usize, + ) -> f64 { + let ks_index = operations_variance + .index + .keyswitch_to_small(src_partition, dst_partition); + let actual_ks_variance = operations_variance.values[ks_index]; + + let mut smallest_ks_max_variance = std::f64::MAX; + + for constraint in &self.undominated_constraints { + let ks_coeff = constraint + .variance + .coeff_keyswitch_to_small(src_partition, dst_partition); + if ks_coeff == 0.0 { + continue; + } + let actual_variance = f64_dot(operations_variance, &constraint.variance.coeffs) + - ks_coeff * actual_ks_variance; + let ks_max_variance = (constraint.safe_variance_bound - actual_variance) / ks_coeff; + smallest_ks_max_variance = smallest_ks_max_variance.min(ks_max_variance); + } + + smallest_ks_max_variance + } + + #[inline(never)] + pub fn fks_max_feasible_variance( + &self, + operations_variance: &OperationsValue, + src_partition: usize, + dst_partition: usize, + ) -> f64 { + let fks_index = operations_variance + .index + .keyswitch_to_big(src_partition, dst_partition); + let actual_fks_variance = operations_variance.values[fks_index]; + + let mut smallest_fks_max_variance = std::f64::MAX; + + for constraint in &self.undominated_constraints { + let fks_coeff = constraint + .variance + .coeff_partition_keyswitch_to_big(src_partition, dst_partition); + if fks_coeff == 0.0 { + continue; + } + let actual_variance = f64_dot(operations_variance, &constraint.variance.coeffs) + - fks_coeff * actual_fks_variance; + let fks_max_variance = (constraint.safe_variance_bound - actual_variance) / fks_coeff; + smallest_fks_max_variance = smallest_fks_max_variance.min(fks_max_variance); + } + + smallest_fks_max_variance + } + + #[inline(never)] pub fn feasible(&self, operations_variance: &OperationsValue) -> bool { if self.global_p_error.is_none() { self.local_feasible(operations_variance) diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/operations_value.rs b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/operations_value.rs index 9251cd8b4c..d3a6620a8f 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/operations_value.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/operations_value.rs @@ -42,18 +42,23 @@ impl Indexing { } pub fn input(self, partition: usize) -> usize { + assert!(partition < self.nb_partitions); partition * self.nb_coeff_per_partition() + VALUE_INDEX_FRESH } pub fn pbs(self, partition: usize) -> usize { + assert!(partition < self.nb_partitions); partition * self.nb_coeff_per_partition() + VALUE_INDEX_PBS } pub fn modulus_switching(self, partition: usize) -> usize { + assert!(partition < self.nb_partitions); partition * self.nb_coeff_per_partition() + VALUE_INDEX_MODULUS } pub fn keyswitch_to_small(self, src_partition: usize, dst_partition: usize) -> usize { + assert!(src_partition < self.nb_partitions); + assert!(dst_partition < self.nb_partitions); // Skip other partition dst_partition * self.nb_coeff_per_partition() // Skip non keyswitchs @@ -63,6 +68,8 @@ impl Indexing { } pub fn keyswitch_to_big(self, src_partition: usize, dst_partition: usize) -> usize { + assert!(src_partition < self.nb_partitions); + assert!(dst_partition < self.nb_partitions); // Skip other partition dst_partition * self.nb_coeff_per_partition() // Skip non keyswitchs diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize.rs b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize.rs index 726c44c3ad..6c9fbb5957 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize.rs @@ -82,14 +82,18 @@ fn optimize_1_ks( cut_complexity: f64, ) -> Option { // find the first feasible (and less complex) + let ks_max_variance = feasible.ks_max_feasible_variance(&operations.variance, ks_src, ks_dst); + let ks_max_cost = complexity.ks_max_cost(cut_complexity, &operations.cost, ks_src, ks_dst); for &ks_quantity in ks_pareto { // variance is decreasing, complexity is increasing - *operations.variance.ks(ks_src, ks_dst) = ks_quantity.noise(ks_input_lwe_dim); - *operations.cost.ks(ks_src, ks_dst) = ks_quantity.complexity(ks_input_lwe_dim); - if complexity.complexity(&operations.cost) > cut_complexity { + let ks_cost = ks_quantity.complexity(ks_input_lwe_dim); + let ks_variance = ks_quantity.noise(ks_input_lwe_dim); + if ks_cost > ks_max_cost { return None; } - if feasible.feasible(&operations.variance) { + if ks_variance <= ks_max_variance { + *operations.variance.ks(ks_src, ks_dst) = ks_variance; + *operations.cost.ks(ks_src, ks_dst) = ks_cost; return Some(ks_quantity); } } @@ -174,6 +178,10 @@ fn optimize_1_fks_and_all_compatible_ks( let mut cut_complexity = cut_complexity; let same_dim = input_glwe == output_glwe; + let fks_max_variance = + feasible.fks_max_feasible_variance(&operations.variance, fks_src, fks_dst); + let mut fks_max_cost = + complexity.fks_max_cost(cut_complexity, &operations.cost, fks_src, fks_dst); for &ks_quantity in &ks_pareto { // OPT: add a pareto cache for fks let fks_quantity = if same_dim { @@ -212,17 +220,20 @@ fn optimize_1_fks_and_all_compatible_ks( dst_glwe_param: output_glwe, } }; - *operations.cost.fks(fks_src, fks_dst) = fks_quantity.complexity; - *operations.variance.fks(fks_src, fks_dst) = fks_quantity.noise; - if complexity.complexity(&operations.cost) > cut_complexity { + if fks_quantity.complexity > fks_max_cost { // complexity is strictly increasing by level // next complexity will be worse return best_sol; } - if !feasible.feasible(&operations.variance) { + + if fks_quantity.noise > fks_max_variance { continue; } + + *operations.cost.fks(fks_src, fks_dst) = fks_quantity.complexity; + *operations.variance.fks(fks_src, fks_dst) = fks_quantity.noise; + let sol = optimize_many_independant_ks( macro_parameters, ks_src, @@ -243,6 +254,7 @@ fn optimize_1_fks_and_all_compatible_ks( continue; } cut_complexity = cost; + fks_max_cost = complexity.fks_max_cost(cut_complexity, &operations.cost, fks_src, fks_dst); // COULD: handle complexity tie let bests = Best1FksAndManyKs { fks: Some((fks_src, fks_quantity)), @@ -334,20 +346,26 @@ fn optimize_1_cmux_and_dst_exclusive_fks_subset_and_all_ks( let mut best_sol_complexity = cut_complexity; let mut best_sol_p_error = best_p_error; let mut best_sol_global_p_error = 1.0; + + let pbs_max_feasible_variance = + feasible.pbs_max_feasible_variance(&operations.variance, partition); for &cmux_quantity in cmux_pareto { // increasing complexity, decreasing variance + + // Lower bounds cuts let pbs_cost = cmux_quantity.complexity_br(internal_dim); *operations.cost.pbs(partition) = pbs_cost; - // Lower bounds cuts let lower_cost = complexity.complexity(&operations.cost); if lower_cost > best_sol_complexity { - continue; + break; // TODO: should be a break ? YES } + let pbs_variance = cmux_quantity.noise_br(internal_dim); - *operations.variance.pbs(partition) = pbs_variance; - if !feasible.feasible(&operations.variance) { + if pbs_variance > pbs_max_feasible_variance { continue; } + + *operations.variance.pbs(partition) = pbs_variance; let sol = optimize_dst_exclusive_fks_subset_and_all_ks( macro_parameters, fks_paretos,