From 4e8b9a199c0ff785424f530f3af7d586a108ff5a Mon Sep 17 00:00:00 2001
From: youben11 <ayoub.benaissa@zama.ai>
Date: Thu, 31 Aug 2023 10:46:39 +0100
Subject: [PATCH] feat(compiler): allow forcing encoding from python

---
 .../lib/Bindings/Python/CompilerAPIModule.cpp       | 11 +++++++++++
 .../Bindings/Python/concrete/compiler/__init__.py   |  2 +-
 .../Python/concrete/compiler/compilation_options.py | 13 +++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp
index 923436b90..e7c9378d8 100644
--- a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp
@@ -64,6 +64,12 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
       .value("DAG_MULTI", optimizer::Strategy::DAG_MULTI)
       .export_values();
 
+  pybind11::enum_<concrete_optimizer::Encoding>(m, "Encoding")
+      .value("AUTO", concrete_optimizer::Encoding::Auto)
+      .value("CRT", concrete_optimizer::Encoding::Crt)
+      .value("NATIVE", concrete_optimizer::Encoding::Native)
+      .export_values();
+
   pybind11::class_<CompilationOptions>(m, "CompilationOptions")
       .def(pybind11::init(
           [](std::string funcname) { return CompilationOptions(funcname); }))
@@ -134,6 +140,11 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
                                     brLevel,       brLogBase,         ksLevel,
                                     ksLogBase,     largeIntegerParam};
            })
+      .def("force_encoding",
+           [](CompilationOptions &options,
+              concrete_optimizer::Encoding encoding) {
+             options.optimizerConfig.encoding = encoding;
+           })
       .def("simulation", [](CompilationOptions &options, bool simulate) {
         options.simulate = simulate;
       });
diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/__init__.py b/compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/__init__.py
index ae1ff6c7d..cffe54b02 100644
--- a/compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/__init__.py
+++ b/compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/__init__.py
@@ -17,7 +17,7 @@ from mlir._mlir_libs._concretelang._compiler import (
 
 # pylint: enable=no-name-in-module,import-error
 
-from .compilation_options import CompilationOptions
+from .compilation_options import CompilationOptions, Encoding
 from .compilation_context import CompilationContext
 from .key_set_cache import KeySetCache
 from .client_parameters import ClientParameters
diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_options.py b/compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_options.py
index 49e04b69c..c25eb56ba 100644
--- a/compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_options.py
+++ b/compilers/concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_options.py
@@ -9,6 +9,7 @@ from typing import List
 from mlir._mlir_libs._concretelang._compiler import (
     CompilationOptions as _CompilationOptions,
     OptimizerStrategy as _OptimizerStrategy,
+    Encoding,
 )
 from .wrapper import WrapperCpp
 
@@ -349,6 +350,18 @@ class CompilationOptions(WrapperCpp):
 
     # pylint: enable=too-many-arguments,too-many-branches
 
+    def force_encoding(self, encoding: Encoding):
+        """Force the compiler to use a specific encoding.
+
+        Args:
+            encoding (Encoding): the encoding to force the compiler to use
+
+        Raises:
+            TypeError: if encoding is not of type Encoding
+        """
+        if not isinstance(encoding, Encoding):
+            raise TypeError("encoding need to be of type Encoding")
+        self.cpp().force_encoding(encoding)
 
     def simulation(self, simulate: bool):
         """Enable or disable simulation.