add ckks support in python frontend

google · Jan 21, 2025 · e4dbe7e · e4dbe7e
1 parent 520071f
commit e4dbe7e
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/heir_py/decorator.py b/heir_py/decorator.py
@@ -134,6 +134,7 @@ def __call__(self, *args, **kwargs):
 
 
 def compile(
+    scheme: str = "bgv",
     backend: str = "openfhe",
     backend_config: Optional[openfhe_config.OpenFHEConfig] = openfhe_config.DEFAULT_INSTALLED_OPENFHE_CONFIG,
     heir_config: Optional[_heir_config.HEIRConfig] = _heir_config.DEVELOPMENT_HEIR_CONFIG,
@@ -156,6 +157,7 @@ def compile(
   def decorator(func):
     compilation_result = run_compiler(
         func,
+        scheme,
         openfhe_config=backend_config or openfhe_config.from_os_env(),
         heir_config=heir_config or _heir_config.from_os_env(),
         debug = debug

diff --git a/heir_py/example.py b/heir_py/example.py
@@ -10,8 +10,8 @@
 # TODO (#1162): Add `scheme` kwarg and add CKKS pipeline (note: addi -> addf)
 # TODO (#1162): Allow manually specifying precision/scale and warn/error if not possible!
 # TODO (#????): Add precision computation/check to Mgmt dialect/infrastructure
-@compile(backend="openfhe", debug=True)
-def foo(x : I16, y : I16):
+@compile(scheme="ckks", backend="openfhe", debug=True)
+def foo(x : F32, y : F32):
   sum = x + y
   diff = x - y
   mul = x * y

diff --git a/heir_py/pipeline.py b/heir_py/pipeline.py
@@ -54,6 +54,7 @@ class CompilationResult:
 
 def run_compiler(
     function,
+    scheme,
     openfhe_config: OpenFHEConfig = openfhe_config_lib.DEFAULT_INSTALLED_OPENFHE_CONFIG,
     heir_config: HEIRConfig = heir_config_lib.DEVELOPMENT_HEIR_CONFIG,
     debug: bool = False,
@@ -106,7 +107,7 @@ def run_compiler(
     heir_opt_options = [
         f"--secretize=function={func_name}",
         (
-            "--mlir-to-openfhe-bgv="
+            f"--mlir-to-openfhe-{scheme}="
             f"entry-function={func_name} ciphertext-degree=32"
         ),
     ]
@@ -129,18 +130,20 @@ def run_compiler(
     pybind_filepath = Path(workspace_dir) / f"{func_name}_bindings.cpp"
     # TODO(#1162): construct heir-translate pipeline options from decorator
     include_type_flag = "--openfhe-include-type=" + openfhe_config.include_type
+    scheme_flag = f"--openfhe-scheme={scheme}"
     heir_translate.run_binary(
         input=heir_opt_output,
         options=[
             "--emit-openfhe-pke-header",
+            scheme_flag,
             include_type_flag,
             "-o",
             h_filepath,
         ],
     )
     heir_translate.run_binary(
         input=heir_opt_output,
-        options=["--emit-openfhe-pke", include_type_flag, "-o", cpp_filepath],
+        options=["--emit-openfhe-pke", scheme_flag, include_type_flag, "-o", cpp_filepath],
     )
     heir_translate.run_binary(
         input=heir_opt_output,