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

@@ -332,12 +332,6 @@ llvm::cl::opt<bool> optimizerNoCacheOnDisk(
                    "cache issues."),
     llvm::cl::init(false));
 
-llvm::cl::opt<bool> optimizerAllowComposition(
-    "optimizer-allow-composition",
-    llvm::cl::desc("Optimizer is parameterized to allow calling the circuit on "
-                   "its own output without decryptions."),
-    llvm::cl::init(false));
-
 llvm::cl::list<int64_t> fhelinalgTileSizes(
     "fhelinalg-tile-sizes",
     llvm::cl::desc(
@@ -508,7 +502,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,23 +396,25 @@ 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(
 func.func @main(%arg0: !FHE.eint<3>) -> !FHE.eint<3> {
-  %cst_1 = arith.constant 1 : i4
-  %1 = "FHE.add_eint_int"(%arg0, %cst_1) : (!FHE.eint<3>, i4) -> !FHE.eint<3>
+  %cst_2 = arith.constant 2 : i4
+  %1 = "FHE.mul_eint_int"(%arg0, %cst_2) : (!FHE.eint<3>, i4) -> !FHE.eint<3>
   return %1: !FHE.eint<3>
 }
 )XXX");
   ASSERT_OUTCOME_HAS_FAILURE_WITH_ERRORMSG(
-      err, "Program can not be composed: No luts in the circuit.");
+      err, "Program can not be composed: Dag is not composable, because of "
+           "output 1: Partition 0 has input coefficient 4");
 }
 
 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);
@@ -430,25 +432,9 @@ func.func @main(%arg0: !FHE.eint<3>) -> (!FHE.eint<3>, !FHE.eint<3>) {
            "output 1: Partition 0 has input coefficient 4");
 }
 
-TEST(CompileComposable, composable_supported_dag_mono) {
-  mlir::concretelang::CompilationOptions options;
-  options.optimizerConfig.composable = true;
-  options.optimizerConfig.display = true;
-  options.optimizerConfig.strategy = mlir::concretelang::optimizer::DAG_MONO;
-  TestProgram circuit(options);
-  auto err = circuit.compile(R"XXX(
-func.func @main(%arg0: !FHE.eint<3>) -> !FHE.eint<3> {
-  %cst_1 = arith.constant 1 : i4
-  %1 = "FHE.add_eint_int"(%arg0, %cst_1) : (!FHE.eint<3>, i4) -> !FHE.eint<3>
-  return %1: !FHE.eint<3>
-}
-)XXX");
-  assert(err.has_value());
-}
-
 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)