From a1398fda9f130cf86aaaded103720f63c79fb466 Mon Sep 17 00:00:00 2001 From: youben11 Date: Wed, 24 Jul 2024 10:42:16 +0100 Subject: [PATCH] feat(optimizer): block macro parameters in external partitions we block them by reducing the search space to a single set of parameters which is provided by the external partition --- .../src/dag/unparametrized.rs | 14 +++ .../dag/multi_parameters/optimize/mod.rs | 25 ++++- .../dag/multi_parameters/optimize/tests.rs | 97 +++++++++++++++++++ .../dag/multi_parameters/partition_cut.rs | 4 +- .../dag/multi_parameters/partitionning.rs | 17 ++++ 5 files changed, 154 insertions(+), 3 deletions(-) diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/dag/unparametrized.rs b/compilers/concrete-optimizer/concrete-optimizer/src/dag/unparametrized.rs index 3fc7aa61c1..3f021756f8 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/dag/unparametrized.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/dag/unparametrized.rs @@ -872,8 +872,20 @@ impl Dag { #[cfg(test)] mod tests { + use crate::{ + optimization::dag::multi_parameters::optimize::MacroParameters, parameters::GlweParameters, + }; + use super::*; + const DUMMY_MACRO_PARAM: MacroParameters = MacroParameters { + glwe_params: GlweParameters { + log2_polynomial_size: 0, + glwe_dimension: 0, + }, + internal_dim: 0, + }; + #[test] fn output_marking() { let mut graph = Dag::new(); @@ -892,6 +904,7 @@ mod tests { let mut graph = Dag::new(); let tfhers_part = ExternalPartition { name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, }; let mut builder = graph.builder("main1"); let a = builder.add_input(1, Shape::number(), Location::Unknown); @@ -942,6 +955,7 @@ mod tests { let tfhers_part = ExternalPartition { name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, }; let change_part = builder.add_change_partition(lut2, Some(&tfhers_part), None, Location::Unknown); diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/mod.rs b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/mod.rs index c146803ab9..a5b2183848 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/mod.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/mod.rs @@ -45,7 +45,7 @@ struct PartialMicroParameters { complexity: f64, } -#[derive(Debug, Copy, Clone, Eq, PartialEq)] +#[derive(Hash, Debug, Copy, Clone, Eq, PartialEq)] pub struct MacroParameters { pub glwe_params: GlweParameters, pub internal_dim: u64, @@ -920,6 +920,11 @@ pub fn optimize( ciphertext_modulus_log, }; + let dag_p_cut = p_cut.as_ref().map_or_else( + || PartitionCut::for_each_precision(dag), + std::clone::Clone::clone, + ); + let dag = analyze(dag, &noise_config, p_cut, default_partition)?; let kappa = error::sigma_scale_of_error_probability(config.maximum_acceptable_error_probability); @@ -954,11 +959,27 @@ pub fn optimize( let mut best_params: Option = None; for iter in 0..=10 { for partition in PartitionIndex::range(0, nb_partitions).rev() { + // reduce search space to the parameters of external partitions + let partition_search_space = if dag_p_cut.is_external_partition(&partition) { + let external_part = + &dag_p_cut.external_partitions[partition.0 - dag_p_cut.n_internal_partitions()]; + let mut reduced_search_space = search_space.clone(); + reduced_search_space.glwe_dimensions = + [external_part.macro_params.glwe_params.glwe_dimension].to_vec(); + reduced_search_space.glwe_log_polynomial_sizes = + [external_part.macro_params.glwe_params.log2_polynomial_size].to_vec(); + reduced_search_space.internal_lwe_dimensions = + [external_part.macro_params.internal_dim].to_vec(); + reduced_search_space + } else { + search_space.clone() + }; + let new_params = optimize_macro( security_level, ciphertext_modulus_log, fft_precision, - search_space, + &partition_search_space, partition, &used_tlu_keyswitch, &used_conversion_keyswitch, diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/tests.rs b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/tests.rs index 97f7545968..7f50d22c19 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/tests.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/tests.rs @@ -1,6 +1,7 @@ #![allow(clippy::float_cmp)] use once_cell::sync::Lazy; +use optimization::dag::multi_parameters::partition_cut::ExternalPartition; use super::*; use crate::computing_cost::cpu::CpuComplexity; @@ -851,3 +852,99 @@ fn test_bug_with_zero_noise() { let sol = optimize(&dag, &None, PartitionIndex(0)); assert!(sol.is_some()); } + +const DUMMY_MACRO_PARAM: MacroParameters = MacroParameters { + glwe_params: GlweParameters { + log2_polynomial_size: 11, + glwe_dimension: 1, + }, + internal_dim: 887, +}; + +#[test] +fn test_optimize_tfhers_in_out_dot_compute() { + let mut dag = unparametrized::Dag::new(); + let input1 = dag.add_input(16, Shape::number()); + let tfhers_partition = ExternalPartition { + name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, + }; + let change_part1 = dag.add_change_partition(input1, Some(&tfhers_partition), None); + let dot = dag.add_dot([change_part1], [2]); + _ = dag.add_change_partition(dot, None, Some(&tfhers_partition)); + + let sol = optimize(&dag, &None, PartitionIndex(0)); + assert!(sol.is_some()); + println!("solution: {:?}", sol.unwrap()); +} + +#[test] +fn test_optimize_tfhers_2lut_compute() { + let mut dag = unparametrized::Dag::new(); + let tfhers_precision = 11; + let input = dag.add_input(tfhers_precision, Shape::number()); + let tfhers_partition_in = ExternalPartition { + name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, + }; + let tfhers_partition_out = ExternalPartition { + name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, + }; + let change_part1 = dag.add_change_partition(input, Some(&tfhers_partition_in), None); + let lut1 = dag.add_lut(change_part1, FunctionTable::UNKWOWN, 4); + let lut2 = dag.add_lut(lut1, FunctionTable::UNKWOWN, tfhers_precision); + let _ = dag.add_change_partition(lut2, None, Some(&tfhers_partition_out)); + + let sol = optimize(&dag, &None, PartitionIndex(0)); + assert!(sol.is_some()); +} + +#[test] +fn test_optimize_tfhers_different_in_out_2lut_compute() { + let mut dag = unparametrized::Dag::new(); + let tfhers_precision = 8; + let input = dag.add_input(tfhers_precision, Shape::number()); + let tfhers_partition_in = ExternalPartition { + name: String::from("tfhers_in"), + macro_params: DUMMY_MACRO_PARAM, + }; + let tfhers_partition_out = ExternalPartition { + name: String::from("tfhers_out"), + macro_params: DUMMY_MACRO_PARAM, + }; + let change_part1 = dag.add_change_partition(input, Some(&tfhers_partition_in), None); + let lut1 = dag.add_lut(change_part1, FunctionTable::UNKWOWN, 4); + let lut2 = dag.add_lut(lut1, FunctionTable::UNKWOWN, tfhers_precision); + let _ = dag.add_change_partition(lut2, None, Some(&tfhers_partition_out)); + + let sol = optimize(&dag, &None, PartitionIndex(0)); + assert!(sol.is_some()); +} + +#[test] +fn test_optimize_tfhers_to_concrete_and_back_example() { + let tfhers_partition = ExternalPartition { + name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, + }; + let concrete_precision = 8; + let msg_width = 2; + let carry_width = 2; + let tfhers_precision = msg_width + carry_width; + + let mut dag = unparametrized::Dag::new(); + let input = dag.add_input( + tfhers_precision, + Shape::vector((concrete_precision / msg_width).into()), + ); + // to concrete + let change_part1 = dag.add_change_partition(input, Some(&tfhers_partition), None); + let lut1 = dag.add_lut(change_part1, FunctionTable::UNKWOWN, concrete_precision); + // from concrete + let lut2 = dag.add_lut(lut1, FunctionTable::UNKWOWN, tfhers_precision); + let _ = dag.add_change_partition(lut2, None, Some(&tfhers_partition)); + + let sol = optimize(&dag, &None, PartitionIndex(0)); + assert!(sol.is_some()); +} diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs index 2cf1eea3f5..8eb11868da 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs @@ -9,13 +9,15 @@ use crate::optimization::dag::multi_parameters::partitions::PartitionIndex; use crate::optimization::dag::solo_key::analyze::out_variances; use crate::optimization::dag::solo_key::symbolic_variance::SymbolicVariance; +use super::optimize::MacroParameters; + const ROUND_INNER_MULTI_PARAMETER: bool = false; const ROUND_EXTERNAL_MULTI_PARAMETER: bool = !ROUND_INNER_MULTI_PARAMETER && true; #[derive(Hash, Eq, PartialEq, Clone, Debug)] pub struct ExternalPartition { pub name: String, - // TODO add params (maybe just macros) + pub macro_params: MacroParameters, } impl std::fmt::Display for ExternalPartition { diff --git a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs index 7905c1ec26..1daf7c25e6 100644 --- a/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs +++ b/compilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs @@ -274,7 +274,9 @@ pub mod tests { use super::*; use crate::dag::operator::{FunctionTable, Shape, Weights}; use crate::dag::unparametrized; + use crate::optimization::dag::multi_parameters::optimize::MacroParameters; use crate::optimization::dag::multi_parameters::partition_cut::ExternalPartition; + use crate::parameters::GlweParameters; fn default_p_cut() -> PartitionCut { PartitionCut::from_precisions(&[2, 128]) @@ -357,12 +359,21 @@ pub mod tests { PartitionIndex(default) } + const DUMMY_MACRO_PARAM: MacroParameters = MacroParameters { + glwe_params: GlweParameters { + log2_polynomial_size: 0, + glwe_dimension: 0, + }, + internal_dim: 0, + }; + #[test] fn test_tfhers_in_out_dot_compute() { let mut dag = unparametrized::Dag::new(); let input1 = dag.add_input(16, Shape::number()); let tfhers_partition = ExternalPartition { name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, }; let change_part1 = dag.add_change_partition(input1, Some(&tfhers_partition), None); let dot = dag.add_dot([change_part1], [2]); @@ -381,6 +392,7 @@ pub mod tests { let input = dag.add_input(16, Shape::number()); let tfhers_partition = ExternalPartition { name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, }; let change_part1 = dag.add_change_partition(input, Some(&tfhers_partition), None); let lut = dag.add_lut(change_part1, FunctionTable::UNKWOWN, 16); @@ -416,9 +428,11 @@ pub mod tests { let input = dag.add_input(16, Shape::number()); let tfhers_partition_in = ExternalPartition { name: String::from("tfhers_in"), + macro_params: DUMMY_MACRO_PARAM, }; let tfhers_partition_out = ExternalPartition { name: String::from("tfhers_out"), + macro_params: DUMMY_MACRO_PARAM, }; let change_part1 = dag.add_change_partition(input, Some(&tfhers_partition_in), None); let lut = dag.add_lut(change_part1, FunctionTable::UNKWOWN, 16); @@ -458,6 +472,7 @@ pub mod tests { let input = dag.add_input(16, Shape::number()); let tfhers_partition = ExternalPartition { name: String::from("tfhers"), + macro_params: DUMMY_MACRO_PARAM, }; let change_part1 = dag.add_change_partition(input, Some(&tfhers_partition), None); let lut1 = dag.add_lut(change_part1, FunctionTable::UNKWOWN, 4); @@ -507,9 +522,11 @@ pub mod tests { let input = dag.add_input(16, Shape::number()); let tfhers_partition_in = ExternalPartition { name: String::from("tfhers_in"), + macro_params: DUMMY_MACRO_PARAM, }; let tfhers_partition_out = ExternalPartition { name: String::from("tfhers_out"), + macro_params: DUMMY_MACRO_PARAM, }; let change_part1 = dag.add_change_partition(input, Some(&tfhers_partition_in), None); let lut1 = dag.add_lut(change_part1, FunctionTable::UNKWOWN, 4);