CKKS Bug Fixes for Relin/Rotate/Mod/Rescale (RNS > 3)

* fixed issues with mod with RNS > 3 for CKKS * fixed issue with rescale for rns > 3 * Added new option in Polys.from_polys to set RNS=1 to support the bug fixes.
IntelLabs · Nov 13, 2024 · 9ed33df · 9ed33df
1 parent dfd7429
commit 9ed33df
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/kerngen/high_parser/types.py b/kerngen/high_parser/types.py
@@ -57,6 +57,9 @@ def from_polys(cls, poly: "Polys", *, mode: str | None = None) -> "Polys":
             case "last_rns":
                 copy.start_rns = copy.rns - 1
                 return cls(**vars(copy))
+            case "single_rns":
+                copy.rns = 1
+                return cls(**vars(copy))
             case "last_part":
                 copy.start_parts = copy.parts - 1
                 return cls(**vars(copy))

diff --git a/kerngen/pisa_generators/mod.py b/kerngen/pisa_generators/mod.py
@@ -123,7 +123,7 @@ def generate_mod_stages() -> list[Stage]:
                                 temp_input_last_rns,
                                 p_half,
                                 Polys.from_polys(
-                                    input_remaining_rns, mode="drop_last_rns"
+                                    input_remaining_rns, mode="single_rns"
                                 ),
                                 last_q,
                             ),

diff --git a/kerngen/pisa_generators/rescale.py b/kerngen/pisa_generators/rescale.py
@@ -62,7 +62,7 @@ def to_pisa(self) -> list[PIsaOp]:
                     temp_input_last_rns,
                     temp_input_last_rns,
                     q_last_half,
-                    input_remaining_rns,
+                    Polys.from_polys(input_remaining_rns, mode="single_rns"),
                     last_q,
                 ),
                 Comment("Subtract q_i (last half/last rns) from y"),