feat(optimizer): support external partition in virtual circuits

compilers/concrete-compiler/compiler/include/concretelang/Common/Keysets.h

@@ -98,8 +98,11 @@ class KeysetCache {
 };
 
 Message<concreteprotocol::KeysetInfo> keysetInfoFromVirtualCircuit(
-    std::vector<concrete_optimizer::utils::PartitionDefinition> partitions,
-    bool generate_fks, std::optional<concrete_optimizer::Options> options);
+    std::vector<concrete_optimizer::utils::InternalPartitionDefinition>
+        internalPartitions,
+    std::vector<concrete_optimizer::utils::ExternalPartitionDefinition>
+        externalPartitions,
+    std::optional<concrete_optimizer::Options> options);
 
 } // namespace keysets
 } // namespace concretelang

compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp

@@ -647,6 +647,12 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
           },
           "Enable or disable overflow detection during simulation.",
           arg("enable_overflow_detection"))
+      .def(
+          "get_optimizer_options",
+          [](CompilationOptions &options) -> concrete_optimizer::Options {
+              return options_from_config(options.optimizerConfig);
+          },
+          "Returns the associated optimizer configuration")
       .doc() = "Holds different flags and options of the compilation process.";
 
   pybind11::enum_<mlir::concretelang::PrimitiveOperation>(m,
@@ -1069,35 +1075,57 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
   // ------------------------------------------------------------------------------//
   // PARTITION DEFINITION //
   // ------------------------------------------------------------------------------//
-  //
-  pybind11::class_<concrete_optimizer::utils::PartitionDefinition>(
-      m, "PartitionDefinition")
+  pybind11::class_<concrete_optimizer::utils::InternalPartitionDefinition>(
+      m, "InternalPartitionDefinition")
       .def(init([](uint8_t precision, double norm2)
-                    -> concrete_optimizer::utils::PartitionDefinition {
-             return concrete_optimizer::utils::PartitionDefinition{precision,
-                                                                   norm2};
+                    -> concrete_optimizer::utils::InternalPartitionDefinition {
+             return concrete_optimizer::utils::InternalPartitionDefinition{
+                 precision, norm2};
            }),
            arg("precision"), arg("norm2"))
       .doc() = "Definition of a partition (in terms of precision in bits and "
                "norm2 in value).";
 
+  // ------------------------------------------------------------------------------//
+  // EXTERNAL PARTITION //
+  // ------------------------------------------------------------------------------//
+  pybind11::class_<concrete_optimizer::utils::ExternalPartitionDefinition>(
+      m, "ExternalPartitionDefinition")
+      .def(init([](std::string name, uint64_t macroLog2PolynomialSize,
+                   uint64_t macroGlweDimension, uint64_t macroInternalDim,
+                   uint64_t pbsLevel, uint64_t pbsBaseLog)
+                    -> concrete_optimizer::utils::ExternalPartitionDefinition {
+             return concrete_optimizer::utils::ExternalPartitionDefinition{
+                 name,
+                 macroLog2PolynomialSize,
+                 macroGlweDimension,
+                 macroInternalDim,
+                 pbsLevel,
+                 pbsBaseLog};
+           }),
+           arg("name"), arg("macro_log2_polynomial_size"),
+           arg("macro_glwe_dimension"), arg("macro_internal_dim"),
+           arg("pbs_level"), arg("pbs_base_log"))
+      .doc() = "Definition of an external partition.";
+
   // ------------------------------------------------------------------------------//
   // KEYSET INFO //
   // ------------------------------------------------------------------------------//
   typedef Message<concreteprotocol::KeysetInfo> KeysetInfo;
   pybind11::class_<KeysetInfo>(m, "KeysetInfo")
       .def_static(
           "generate_virtual",
-          [](std::vector<concrete_optimizer::utils::PartitionDefinition>
-                 partitions,
-             bool generateFks,
+          [](std::vector<concrete_optimizer::utils::InternalPartitionDefinition>
+                 internalPartitions,
+             std::vector<concrete_optimizer::utils::ExternalPartitionDefinition>
+                 externalPartitions,
              std::optional<concrete_optimizer::Options> options) -> KeysetInfo {
-            if (partitions.size() < 2) {
+            if (internalPartitions.size() + externalPartitions.size() < 2) {
               throw std::runtime_error("Need at least two partition defs to "
                                        "generate a virtual keyset info.");
             }
             return ::concretelang::keysets::keysetInfoFromVirtualCircuit(
-                partitions, generateFks, options);
+                internalPartitions, externalPartitions, options);
           },
           arg("partition_defs"), arg("generate_fks"),
           arg("options") = std::nullopt,

compilers/concrete-compiler/compiler/lib/Common/Keysets.cpp

@@ -559,14 +559,22 @@ generateKeysetInfoFromParameters(CircuitKeys parameters,
 }
 
 Message<concreteprotocol::KeysetInfo> keysetInfoFromVirtualCircuit(
-    std::vector<concrete_optimizer::utils::PartitionDefinition> partitionDefs,
-    bool generateFks, std::optional<concrete_optimizer::Options> options) {
+    std::vector<concrete_optimizer::utils::InternalPartitionDefinition> partitionDefs,
+    std::vector<concrete_optimizer::utils::ExternalPartitionDefinition>
+        externalPartitions,
+    std::optional<concrete_optimizer::Options> options) {
 
-  rust::Vec<concrete_optimizer::utils::PartitionDefinition> rustPartitionDefs{};
+  rust::Vec<concrete_optimizer::utils::InternalPartitionDefinition> rustPartitionDefs{};
   for (auto def : partitionDefs) {
     rustPartitionDefs.push_back(def);
   }
 
+  rust::Vec<concrete_optimizer::utils::ExternalPartitionDefinition>
+      rustExternalPartitions{};
+  for (auto def : externalPartitions) {
+    rustExternalPartitions.push_back(def);
+  }
+
   auto defaultOptions = concrete_optimizer::Options{};
   defaultOptions.security_level = 128;
   defaultOptions.maximum_acceptable_error_probability = 0.000063342483999973;
@@ -577,7 +585,7 @@ Message<concreteprotocol::KeysetInfo> keysetInfoFromVirtualCircuit(
   auto opts = options.value_or(defaultOptions);
 
   auto parameters = concrete_optimizer::utils::generate_virtual_keyset_info(
-      rustPartitionDefs, generateFks, opts);
+      rustPartitionDefs, rustExternalPartitions, opts);
 
   return generateKeysetInfoFromParameters(parameters, opts);
 }

compilers/concrete-compiler/compiler/lib/Dialect/FHE/Analysis/ConcreteOptimizer.cpp

@@ -274,8 +274,8 @@ struct FunctionToDag {
               options, logPolySize, glweDim, lweDim, pbsLevel, pbsLogBase);
           auto name = partitionAttr.getName().getValue().str();
           // TODO: max_variance vs variance
-          return concrete_optimizer::utils::get_external_partition(
-              name, logPolySize, glweDim, lweDim, max_variance, max_variance);
+          return concrete_optimizer::utils::ExternalPartition{
+              name, logPolySize, glweDim, lweDim, max_variance, max_variance};
         };
 
     if (auto srcPartitionAttr =
@@ -287,14 +287,14 @@ struct FunctionToDag {
       auto partition = partitionBuilder(srcPartitionAttr);
 
       index[val] = dagBuilder.add_change_partition_with_src(
-          encrypted_input, *partition, *loc_to_location(val.getLoc()));
+          encrypted_input, partition, *loc_to_location(val.getLoc()));
     } else if (auto destPartitionAttr =
                    op.getAttrOfType<mlir::concretelang::FHE::PartitionAttr>(
                        "dest")) {
       auto partition = partitionBuilder(destPartitionAttr);
 
       index[val] = dagBuilder.add_change_partition_with_dst(
-          encrypted_input, *partition, *loc_to_location(val.getLoc()));
+          encrypted_input, partition, *loc_to_location(val.getLoc()));
     } else {
       assert(false &&
              "ChangePartition: one of src or dest partitions need to be set");

compilers/concrete-compiler/compiler/lib/Support/V0Parameters.cpp

@@ -391,6 +391,8 @@ getSolution(optimizer::Description &descr, ProgramCompilationFeedback &feedback,
       if (sol.is_feasible || !config.composition_rules.empty()) {
         displayOptimizer(sol, descr, config);
         return toCompilerSolution(sol, feedback, config);
+      } else {
+          assert(false);
       }
     }
     config.strategy = optimizer::Strategy::DAG_MONO;

compilers/concrete-optimizer/concrete-optimizer-cpp/src/concrete-optimizer.rs

@@ -1,7 +1,7 @@
-#![allow(clippy::boxed_local)]
-#![allow(clippy::too_many_arguments)]
+// #![allow(clippy::boxed_local)]
+// #![allow(clippy::too_many_arguments)]
 
-use core::panic;
+// use core::panic;
 
 use concrete_optimizer::computing_cost::cpu::CpuComplexity;
 use concrete_optimizer::config;
@@ -16,7 +16,9 @@ use concrete_optimizer::optimization::dag::multi_parameters::optimize::{
     KeysetRestriction, MacroParameters, NoSearchSpaceRestriction, RangeRestriction,
     SearchSpaceRestriction,
 };
-use concrete_optimizer::optimization::dag::multi_parameters::partition_cut::PartitionCut;
+use concrete_optimizer::optimization::dag::multi_parameters::partition_cut::{
+    ExternalPartition, PartitionCut,
+};
 use concrete_optimizer::optimization::dag::multi_parameters::virtual_circuit::generate_virtual_parameters;
 use concrete_optimizer::optimization::dag::multi_parameters::{keys_spec, PartitionIndex};
 use concrete_optimizer::optimization::dag::solo_key::optimize_generic::{
@@ -62,35 +64,6 @@ fn caches_from(options: &ffi::Options) -> decomposition::PersistDecompCaches {
     )
 }
 
-#[derive(Clone)]
-pub struct ExternalPartition(
-    concrete_optimizer::optimization::dag::multi_parameters::partition_cut::ExternalPartition,
-);
-
-pub fn get_external_partition(
-    name: String,
-    log2_polynomial_size: u64,
-    glwe_dimension: u64,
-    internal_dim: u64,
-    max_variance: f64,
-    variance: f64,
-) -> Box<ExternalPartition> {
-    Box::new(ExternalPartition(
-        concrete_optimizer::optimization::dag::multi_parameters::partition_cut::ExternalPartition {
-            name,
-            macro_params: MacroParameters {
-                glwe_params: GlweParameters {
-                    log2_polynomial_size,
-                    glwe_dimension,
-                },
-                internal_dim,
-            },
-            max_variance,
-            variance,
-        },
-    ))
-}
-
 pub fn get_noise_br(
     options: &ffi::Options,
     log2_polynomial_size: u64,
@@ -852,13 +825,13 @@ impl DagBuilder<'_> {
     fn add_change_partition_with_src(
         &mut self,
         input: ffi::OperatorIndex,
-        src_partition: &ExternalPartition,
+        src_partition: &ffi::ExternalPartition,
         location: &Location,
     ) -> ffi::OperatorIndex {
         self.0
             .add_change_partition(
                 input.into(),
-                Some(src_partition.0.clone()),
+                Some(src_partition.clone().into()),
                 None,
                 location.0.clone(),
             )
@@ -868,14 +841,14 @@ impl DagBuilder<'_> {
     fn add_change_partition_with_dst(
         &mut self,
         input: ffi::OperatorIndex,
-        dst_partition: &ExternalPartition,
+        dst_partition: &ffi::ExternalPartition,
         location: &Location,
     ) -> ffi::OperatorIndex {
         self.0
             .add_change_partition(
                 input.into(),
                 None,
-                Some(dst_partition.0.clone()),
+                Some(dst_partition.clone().into()),
                 location.0.clone(),
             )
             .into()
@@ -915,8 +888,8 @@ fn location_from_string(string: &str) -> Box<Location> {
 }
 
 fn generate_virtual_keyset_info(
-    inputs: Vec<ffi::PartitionDefinition>,
-    generate_fks: bool,
+    internal_partitions: Vec<ffi::InternalPartitionDefinition>,
+    external_partitions: Vec<ffi::ExternalPartitionDefinition>,
     options: &ffi::Options,
 ) -> ffi::CircuitKeys {
     let config = Config {
@@ -927,14 +900,18 @@ fn generate_virtual_keyset_info(
         fft_precision: options.fft_precision,
         complexity_model: &CpuComplexity::default(),
     };
+
     generate_virtual_parameters(
-        inputs
+        internal_partitions
             .into_iter()
             .map(
-                |ffi::PartitionDefinition { precision, norm2 }| concrete_optimizer::optimization::dag::multi_parameters::virtual_circuit::PartitionDefinition { precision, norm2 },
+                |ffi::InternalPartitionDefinition { precision, norm2 }| concrete_optimizer::optimization::dag::multi_parameters::virtual_circuit::InternalPartition { precision, norm2 },
             )
             .collect(),
-        generate_fks,
+        external_partitions.into_iter().map(|def| {
+            let noise = get_noise_br(options, def.log2_polynomial_size, def.glwe_dimension, def.internal_dim, def.pbs_level, def.pbs_base_log);
+            ExternalPartition { name: def.name, macro_params: MacroParameters { glwe_params: GlweParameters { log2_polynomial_size: def.log2_polynomial_size, glwe_dimension: def.glwe_dimension }, internal_dim: def.internal_dim } , max_variance: noise, variance: noise }
+        }).collect(),
         config
     )
     .into()
@@ -975,6 +952,24 @@ impl Into<Encoding> for ffi::Encoding {
     }
 }
 
+#[allow(clippy::from_over_into)]
+impl Into<ExternalPartition> for ffi::ExternalPartition {
+    fn into(self) -> ExternalPartition {
+        ExternalPartition {
+            name: self.name,
+            macro_params: MacroParameters {
+                glwe_params: GlweParameters {
+                    log2_polynomial_size: self.macro_log2_polynomial_size,
+                    glwe_dimension: self.macro_glwe_dimension,
+                },
+                internal_dim: self.macro_internal_dim,
+            },
+            max_variance: self.max_variance,
+            variance: self.variance,
+        }
+    }
+}
+
 #[allow(
     unused_must_use,
     clippy::needless_lifetimes,
@@ -1000,8 +995,6 @@ mod ffi {
 
         type Location;
 
-        type ExternalPartition;
-
         #[namespace = "concrete_optimizer::utils"]
         fn location_unknown() -> Box<Location>;
 
@@ -1010,21 +1003,11 @@ mod ffi {
 
         #[namespace = "concrete_optimizer::utils"]
         fn generate_virtual_keyset_info(
-            partitions: Vec<PartitionDefinition>,
-            generate_fks: bool,
+            internal_partitions: Vec<InternalPartitionDefinition>,
+            external_partitions: Vec<ExternalPartitionDefinition>,
             options: &Options,
         ) -> CircuitKeys;
 
-        #[namespace = "concrete_optimizer::utils"]
-        fn get_external_partition(
-            name: String,
-            log2_polynomial_size: u64,
-            glwe_dimension: u64,
-            internal_dim: u64,
-            max_variance: f64,
-            variance: f64,
-        ) -> Box<ExternalPartition>;
-
         #[namespace = "concrete_optimizer::utils"]
         fn get_noise_br(
             options: &Options,
@@ -1396,10 +1379,32 @@ mod ffi {
 
     #[namespace = "concrete_optimizer::utils"]
     #[derive(Debug, Clone)]
-    pub struct PartitionDefinition {
+    pub struct InternalPartitionDefinition {
         pub precision: u8,
         pub norm2: f64,
     }
+
+    #[namespace = "concrete_optimizer::utils"]
+    #[derive(Debug, Clone)]
+    pub struct ExternalPartitionDefinition {
+        pub name: String,
+        pub log2_polynomial_size: u64,
+        pub glwe_dimension: u64,
+        pub internal_dim: u64,
+        pub pbs_level: u64,
+        pub pbs_base_log: u64,
+    }
+
+    #[namespace = "concrete_optimizer::utils"]
+    #[derive(Debug, Clone)]
+    pub struct ExternalPartition {
+        pub name: String,
+        pub macro_log2_polynomial_size: u64,
+        pub macro_glwe_dimension: u64,
+        pub macro_internal_dim: u64,
+        pub max_variance: f64,
+        pub variance: f64,
+    }
 }
 
 fn processing_unit(options: &ffi::Options) -> ProcessingUnit {