feat(frontend): add support for wires to concrete-python

compilers/concrete-compiler/compiler/include/concretelang/Support/V0Parameters.h

@@ -67,20 +67,19 @@ struct V0Parameter {
 };
 
 namespace optimizer {
-constexpr double DEFAULT_GLOBAL_P_ERROR = 1.0 / 100000.0;
-constexpr double UNSPECIFIED_P_ERROR = NAN; // will use DEFAULT_GLOBAL_P_ERROR
-constexpr double UNSPECIFIED_GLOBAL_P_ERROR =
+const double DEFAULT_GLOBAL_P_ERROR = 1.0 / 100000.0;
+const double UNSPECIFIED_P_ERROR = NAN; // will use DEFAULT_GLOBAL_P_ERROR
+const double UNSPECIFIED_GLOBAL_P_ERROR =
     NAN; // will use DEFAULT_GLOBAL_P_ERROR
-constexpr uint DEFAULT_SECURITY = 128;
-constexpr uint DEFAULT_FALLBACK_LOG_NORM_WOPPBS = 8;
-constexpr bool DEFAULT_DISPLAY = false;
-constexpr bool DEFAULT_USE_GPU_CONSTRAINTS = false;
-constexpr concrete_optimizer::Encoding DEFAULT_ENCODING =
+const uint DEFAULT_SECURITY = 128;
+const uint DEFAULT_FALLBACK_LOG_NORM_WOPPBS = 8;
+const bool DEFAULT_DISPLAY = false;
+const bool DEFAULT_USE_GPU_CONSTRAINTS = false;
+const concrete_optimizer::Encoding DEFAULT_ENCODING =
     concrete_optimizer::Encoding::Auto;
-constexpr bool DEFAULT_CACHE_ON_DISK = true;
-constexpr uint32_t DEFAULT_CIPHERTEXT_MODULUS_LOG = 64;
-constexpr uint32_t DEFAULT_FFT_PRECISION = 53;
-constexpr bool DEFAULT_COMPOSABLE = false;
+const bool DEFAULT_CACHE_ON_DISK = true;
+const uint32_t DEFAULT_CIPHERTEXT_MODULUS_LOG = 64;
+const uint32_t DEFAULT_FFT_PRECISION = 53;
 
 /// The strategy of the crypto optimization
 enum Strategy {
@@ -96,10 +95,19 @@ enum Strategy {
 
 std::string const StrategyLabel[] = {"V0", "dag-mono", "dag-multi"};
 
-constexpr Strategy DEFAULT_STRATEGY = Strategy::DAG_MULTI;
-constexpr concrete_optimizer::MultiParamStrategy DEFAULT_MULTI_PARAM_STRATEGY =
+const Strategy DEFAULT_STRATEGY = Strategy::DAG_MULTI;
+const concrete_optimizer::MultiParamStrategy DEFAULT_MULTI_PARAM_STRATEGY =
     concrete_optimizer::MultiParamStrategy::ByPrecision;
-constexpr bool DEFAULT_KEY_SHARING = true;
+const bool DEFAULT_KEY_SHARING = true;
+
+struct CompositionRule {
+  std::string from_func;
+  size_t from_pos;
+  std::string to_func;
+  size_t to_pos;
+};
+
+const std::vector<CompositionRule> DEFAULT_COMPOSITION_RULES = {};
 
 struct Config {
   double p_error;
@@ -115,10 +123,10 @@ struct Config {
   bool cache_on_disk;
   uint32_t ciphertext_modulus_log;
   uint32_t fft_precision;
-  bool composable;
+  std::vector<CompositionRule> composition_rules;
 };
 
-constexpr Config DEFAULT_CONFIG = {
+const Config DEFAULT_CONFIG = {
     UNSPECIFIED_P_ERROR,
     UNSPECIFIED_GLOBAL_P_ERROR,
     DEFAULT_DISPLAY,
@@ -132,7 +140,7 @@ constexpr Config DEFAULT_CONFIG = {
     DEFAULT_CACHE_ON_DISK,
     DEFAULT_CIPHERTEXT_MODULUS_LOG,
     DEFAULT_FFT_PRECISION,
-    DEFAULT_COMPOSABLE,
+    DEFAULT_COMPOSITION_RULES,
 };
 
 using Dag = rust::Box<concrete_optimizer::Dag>;

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

@@ -4,6 +4,7 @@
 // for license information.
 
 #include "concretelang/Bindings/Python/CompilerAPIModule.h"
+#include "concrete-optimizer.hpp"
 #include "concrete-protocol.capnp.h"
 #include "concretelang/ClientLib/ClientLib.h"
 #include "concretelang/Common/Compat.h"
@@ -744,9 +745,11 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
            [](CompilationOptions &options, double global_p_error) {
              options.optimizerConfig.global_p_error = global_p_error;
            })
-      .def("set_composable",
-           [](CompilationOptions &options, bool composable) {
-             options.optimizerConfig.composable = composable;
+      .def("add_composition",
+           [](CompilationOptions &options, std::string from_func,
+              size_t from_pos, std::string to_func, size_t to_pos) {
+             options.optimizerConfig.composition_rules.push_back(
+                 {from_func, from_pos, to_func, to_pos});
            })
       .def("set_security_level",
            [](CompilationOptions &options, int security_level) {

compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_options.py

@@ -61,18 +61,27 @@ def new(backend=_Backend.CPU) -> "CompilationOptions":
 
     # pylint: enable=arguments-differ
 
-    def set_composable(self, composable: bool):
-        """Set option for composition.
+    def add_composition(self, from_func: str, from_pos: int, to_func: str, to_pos: int):
+        """Adds a composition rule.
 
         Args:
-            composable (bool): whether to turn it on or off
+            from_func(str): the name of the circuit the output comes from.
+            from_pos(int): the return position of the output.
+            to_func(str): the name of the circuit the input targets.
+            to_pos(int): the argument position of the input.
 
         Raises:
-            TypeError: if the value to set is not boolean
+            TypeError: if the inputs do not have the proper type.
         """
-        if not isinstance(composable, bool):
-            raise TypeError("can't set the option to a non-boolean value")
-        self.cpp().set_composable(composable)
+        if not isinstance(from_func, str):
+            raise TypeError("expected `from_func` to be (str)")
+        if not isinstance(from_pos, int):
+            raise TypeError("expected `from_pos` to be (int)")
+        if not isinstance(to_func, str):
+            raise TypeError("expected `to_func` to be (str)")
+        if not isinstance(from_pos, int):
+            raise TypeError("expected `to_pos` to be (int)")
+        self.cpp().add_composition(from_func, from_pos, to_func, to_pos)
 
     def set_auto_parallelize(self, auto_parallelize: bool):
         """Set option for auto parallelization.

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

@@ -976,15 +976,9 @@ std::unique_ptr<mlir::Pass> createDagPass(optimizer::Config config,
 // Adds the composition rules to the
 void applyCompositionRules(optimizer::Config config,
                            concrete_optimizer::Dag &dag) {
-
-  if (config.composable) {
-    auto inputs = dag.get_input_indices();
-    auto outputs = dag.get_output_indices();
-    dag.add_compositions(
-        rust::Slice<const concrete_optimizer::dag::OperatorIndex>(
-            outputs.data(), outputs.size()),
-        rust::Slice<const concrete_optimizer::dag::OperatorIndex>(
-            inputs.data(), inputs.size()));
+  for (auto rule : config.composition_rules) {
+    dag.add_composition(rule.from_func, rule.from_pos, rule.to_func,
+                        rule.to_pos);
   }
 }
 

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

@@ -178,12 +178,6 @@ CompilerEngine::getConcreteOptimizerDescription(CompilationResult &res) {
   if (!description->has_value()) { // The pass has not been run
     return std::nullopt;
   }
-  if (description->value().dag.value()->get_circuit_count() > 1 &&
-      config.strategy !=
-          mlir::concretelang::optimizer::V0) { // Multi circuits without V0
-    return StreamStringError(
-        "Multi-circuits is only supported for V0 optimization.");
-  }
   return description;
 }
 

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

@@ -377,7 +377,7 @@ getSolution(optimizer::Description &descr, ProgramCompilationFeedback &feedback,
     if (encoding != concrete_optimizer::Encoding::Crt) {
       config.encoding = concrete_optimizer::Encoding::Native;
       auto sol = getDagMultiSolution(descr.dag.value(), config);
-      if (sol.is_feasible || config.composable) {
+      if (sol.is_feasible || !config.composition_rules.empty()) {
         displayOptimizer(sol, descr, config);
         return toCompilerSolution(sol, feedback, config);
       }

compilers/concrete-compiler/compiler/src/main.cpp

@@ -508,7 +508,6 @@ cmdlineCompilationOptions() {
       cmdline::optimizerMultiParamStrategy;
   options.optimizerConfig.encoding = cmdline::optimizerEncoding;
   options.optimizerConfig.cache_on_disk = !cmdline::optimizerNoCacheOnDisk;
-  options.optimizerConfig.composable = cmdline::optimizerAllowComposition;
 
   if (!std::isnan(options.optimizerConfig.global_p_error) &&
       options.optimizerConfig.strategy == optimizer::Strategy::V0) {

compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_jit_distributed.cc

@@ -81,7 +81,7 @@ func.func @main(%arg0: tensor<200x4x!FHE.eint<4>>) -> tensor<200x8x!FHE.eint<4>>
 )XXX",
              "main", false, true, true, DEFAULT_batchTFHEOps,
              DEFAULT_global_p_error, DEFAULT_chunkedIntegers, DEFAULT_chunkSize,
-             DEFAULT_chunkWidth, DEFAULT_composable, false);
+             DEFAULT_chunkWidth, false);
 
   const size_t dim0 = 200;
   const size_t dim1 = 4;
@@ -121,8 +121,7 @@ TEST(Distributed, nn_med_sequential) {
 )XXX",
                "main", false, false, false, DEFAULT_batchTFHEOps,
                DEFAULT_global_p_error, DEFAULT_chunkedIntegers,
-               DEFAULT_chunkSize, DEFAULT_chunkWidth, DEFAULT_composable,
-               false);
+               DEFAULT_chunkSize, DEFAULT_chunkWidth, false);
 
     const size_t dim0 = 200;
     const size_t dim1 = 4;

compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_jit_test.cc

@@ -396,7 +396,7 @@ func.func @main(%arg0: !FHE.eint<3>) -> !FHE.eint<3> {
 
 TEST(CompileNotComposable, not_composable_1) {
   mlir::concretelang::CompilationOptions options;
-  options.optimizerConfig.composable = true;
+  options.optimizerConfig.composition_rules.push_back({"main", 0, "main", 0});
   options.optimizerConfig.strategy = mlir::concretelang::optimizer::DAG_MULTI;
   TestProgram circuit(options);
   auto err = circuit.compile(R"XXX(
@@ -412,7 +412,8 @@ func.func @main(%arg0: !FHE.eint<3>) -> !FHE.eint<3> {
 
 TEST(CompileNotComposable, not_composable_2) {
   mlir::concretelang::CompilationOptions options;
-  options.optimizerConfig.composable = true;
+  options.optimizerConfig.composition_rules.push_back({"main", 0, "main", 0});
+  options.optimizerConfig.composition_rules.push_back({"main", 1, "main", 0});
   options.optimizerConfig.display = true;
   options.optimizerConfig.strategy = mlir::concretelang::optimizer::DAG_MULTI;
   TestProgram circuit(options);
@@ -432,7 +433,7 @@ func.func @main(%arg0: !FHE.eint<3>) -> (!FHE.eint<3>, !FHE.eint<3>) {
 
 TEST(CompileComposable, composable_supported_dag_mono) {
   mlir::concretelang::CompilationOptions options;
-  options.optimizerConfig.composable = true;
+  options.optimizerConfig.composition_rules.push_back({"main", 0, "main", 0});
   options.optimizerConfig.display = true;
   options.optimizerConfig.strategy = mlir::concretelang::optimizer::DAG_MONO;
   TestProgram circuit(options);
@@ -448,7 +449,7 @@ func.func @main(%arg0: !FHE.eint<3>) -> !FHE.eint<3> {
 
 TEST(CompileComposable, composable_supported_v0) {
   mlir::concretelang::CompilationOptions options;
-  options.optimizerConfig.composable = true;
+  options.optimizerConfig.composition_rules.push_back({"main", 0, "main", 0});
   options.optimizerConfig.display = true;
   options.optimizerConfig.strategy = mlir::concretelang::optimizer::V0;
   TestProgram circuit(options);

compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_jit_test.h

@@ -25,7 +25,6 @@ double DEFAULT_global_p_error = TEST_ERROR_RATE;
 bool DEFAULT_chunkedIntegers = false;
 unsigned int DEFAULT_chunkSize = 4;
 unsigned int DEFAULT_chunkWidth = 2;
-bool DEFAULT_composable = false;
 bool DEFAULT_use_multi_parameter = true;
 
 // Jit-compiles the function specified by `func` from `src` and
@@ -41,7 +40,6 @@ inline Result<TestProgram> internalCheckedJit(
     bool chunkedIntegers = DEFAULT_chunkedIntegers,
     unsigned int chunkSize = DEFAULT_chunkSize,
     unsigned int chunkWidth = DEFAULT_chunkWidth,
-    bool composable = DEFAULT_composable,
     bool use_multi_parameter = DEFAULT_use_multi_parameter) {
 
   auto options = mlir::concretelang::CompilationOptions();
@@ -59,11 +57,6 @@ inline Result<TestProgram> internalCheckedJit(
   options.dataflowParallelize = dataflowParallelize;
 #endif
   options.batchTFHEOps = batchTFHEOps;
-  if (composable) {
-    options.optimizerConfig.composable = composable;
-    options.optimizerConfig.strategy =
-        mlir::concretelang::optimizer::Strategy::DAG_MULTI;
-  }
 
   if (!use_multi_parameter)
     options.optimizerConfig.strategy =

compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.h

@@ -237,9 +237,6 @@ std::string getOptionsName(mlir::concretelang::CompilationOptions compilation) {
   if (compilation.optimizerConfig.security !=
       defaultOptions.optimizerConfig.security)
     os << "_optimizerSecurity" << compilation.optimizerConfig.security;
-  if (compilation.optimizerConfig.composable !=
-      defaultOptions.optimizerConfig.composable)
-    os << "_optimizerSecurity" << compilation.optimizerConfig.composable;
 
   /// GPU
   if (compilation.emitGPUOps != defaultOptions.emitGPUOps)

compilers/concrete-optimizer/concrete-optimizer-cpp/Makefile

@@ -20,7 +20,7 @@ INTERFACE_CPP =  src/cpp/concrete-optimizer.cpp
 SOURCES = $(shell find $(ROOT)/concrete-optimizer/src) \
 	$(shell find $(ROOT)/concrete-optimizer-cpp/src -name '*.rs')
 
-build: $(INTERFACE_LIB)
+build: $(INTERFACE_LIB) $(INTERFACE_CPP) $(INTERFACE_HEADER)
 
 $(INTERFACE_LIB_ORIG) $(INTERFACE_HEADER_ORIG) $(INTERFACE_CPP_ORIG): $(SOURCES)
 	cd $(ROOT) && cargo build -p concrete-optimizer-cpp --profile $(CARGO_PROFILE)

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

@@ -16,6 +16,7 @@ use concrete_optimizer::optimization::decomposition;
 use concrete_optimizer::parameters::{BrDecompositionParameters, KsDecompositionParameters};
 use concrete_optimizer::utils::cache::persistent::default_cache_dir;
 use concrete_optimizer::utils::viz::Viz;
+use cxx::CxxString;
 
 fn no_solution() -> ffi::Solution {
     ffi::Solution {
@@ -547,16 +548,28 @@ impl Dag {
         self.0.get_circuit_count()
     }
 
-    fn add_compositions(&mut self, froms: &[ffi::OperatorIndex], tos: &[ffi::OperatorIndex]) {
-        self.0.add_compositions(
-            froms
-                .iter()
-                .map(|a| OperatorIndex(a.index))
-                .collect::<Vec<_>>(),
-            tos.iter()
-                .map(|a| OperatorIndex(a.index))
-                .collect::<Vec<_>>(),
-        );
+    unsafe fn add_composition<'a>(
+        &mut self,
+        from_func: &'a CxxString,
+        from_pos: usize,
+        to_func: &'a CxxString,
+        to_pos: usize,
+    ) {
+        let from_index = self
+            .0
+            .get_circuit(from_func.to_str().unwrap())
+            .get_output_operators_iter()
+            .nth(from_pos)
+            .unwrap()
+            .id;
+        let to_index = self
+            .0
+            .get_circuit(to_func.to_str().unwrap())
+            .get_input_operators_iter()
+            .nth(to_pos)
+            .unwrap()
+            .id;
+        self.0.add_composition(from_index, to_index);
     }
 
     fn add_all_compositions(&mut self) {
@@ -800,7 +813,13 @@ mod ffi {
 
         fn optimize(self: &Dag, options: Options) -> DagSolution;
 
-        fn add_compositions(self: &mut Dag, froms: &[OperatorIndex], tos: &[OperatorIndex]);
+        unsafe fn add_composition<'a>(
+            self: &mut Dag,
+            from_func: &'a CxxString,
+            from_pos: usize,
+            to_func: &'a CxxString,
+            to_pos: usize,
+        );
 
         fn add_all_compositions(self: &mut Dag);
 

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

@@ -972,7 +972,7 @@ struct Dag final : public ::rust::Opaque {
   ::rust::Box<::concrete_optimizer::DagBuilder> builder(::rust::String circuit) noexcept;
   ::rust::String dump() const noexcept;
   ::concrete_optimizer::dag::DagSolution optimize(::concrete_optimizer::Options options) const noexcept;
-  void add_compositions(::rust::Slice<::concrete_optimizer::dag::OperatorIndex const> froms, ::rust::Slice<::concrete_optimizer::dag::OperatorIndex const> tos) noexcept;
+  void add_composition(::std::string const &from_func, ::std::size_t from_pos, ::std::string const &to_func, ::std::size_t to_pos) noexcept;
   void add_all_compositions() noexcept;
   ::std::size_t get_circuit_count() const noexcept;
   ::concrete_optimizer::dag::CircuitSolution optimize_multi(::concrete_optimizer::Options options) const noexcept;
@@ -1319,7 +1319,7 @@ void concrete_optimizer$cxxbridge1$DagBuilder$tag_operator_as_output(::concrete_
 
 void concrete_optimizer$cxxbridge1$Dag$optimize(::concrete_optimizer::Dag const &self, ::concrete_optimizer::Options options, ::concrete_optimizer::dag::DagSolution *return$) noexcept;
 
-void concrete_optimizer$cxxbridge1$Dag$add_compositions(::concrete_optimizer::Dag &self, ::rust::Slice<::concrete_optimizer::dag::OperatorIndex const> froms, ::rust::Slice<::concrete_optimizer::dag::OperatorIndex const> tos) noexcept;
+void concrete_optimizer$cxxbridge1$Dag$add_composition(::concrete_optimizer::Dag &self, ::std::string const &from_func, ::std::size_t from_pos, ::std::string const &to_func, ::std::size_t to_pos) noexcept;
 
 void concrete_optimizer$cxxbridge1$Dag$add_all_compositions(::concrete_optimizer::Dag &self) noexcept;
 } // extern "C"
@@ -1449,8 +1449,8 @@ ::concrete_optimizer::dag::DagSolution Dag::optimize(::concrete_optimizer::Optio
   return ::std::move(return$.value);
 }
 
-void Dag::add_compositions(::rust::Slice<::concrete_optimizer::dag::OperatorIndex const> froms, ::rust::Slice<::concrete_optimizer::dag::OperatorIndex const> tos) noexcept {
-  concrete_optimizer$cxxbridge1$Dag$add_compositions(*this, froms, tos);
+void Dag::add_composition(::std::string const &from_func, ::std::size_t from_pos, ::std::string const &to_func, ::std::size_t to_pos) noexcept {
+  concrete_optimizer$cxxbridge1$Dag$add_composition(*this, from_func, from_pos, to_func, to_pos);
 }
 
 void Dag::add_all_compositions() noexcept {

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

@@ -953,7 +953,7 @@ struct Dag final : public ::rust::Opaque {
   ::rust::Box<::concrete_optimizer::DagBuilder> builder(::rust::String circuit) noexcept;
   ::rust::String dump() const noexcept;
   ::concrete_optimizer::dag::DagSolution optimize(::concrete_optimizer::Options options) const noexcept;
-  void add_compositions(::rust::Slice<::concrete_optimizer::dag::OperatorIndex const> froms, ::rust::Slice<::concrete_optimizer::dag::OperatorIndex const> tos) noexcept;
+  void add_composition(::std::string const &from_func, ::std::size_t from_pos, ::std::string const &to_func, ::std::size_t to_pos) noexcept;
   void add_all_compositions() noexcept;
   ::std::size_t get_circuit_count() const noexcept;
   ::concrete_optimizer::dag::CircuitSolution optimize_multi(::concrete_optimizer::Options options) const noexcept;