ITADN

CKKS: basis conversion and modulus switching

#2987Openbencemali 创建于 2026-05-26
dialect: lwe
B
bencemalicommented
While having a look at the [CKKS key switch lowering](https://github.com/google/heir/blob/main/lib/Dialect/CKKS/Transforms/DecomposeKeySwitch.cpp#L34), I noticed that the last step (Step 4: Remove the key-switch primes) incorrectly performs an [LWE level basis conversion op](https://github.com/google/heir/blob/main/lib/Dialect/LWE/IR/LWEOps.td#L324) from `QP` to `Q` instead of a modulus switch from `QP` to `Q`. These operations would correspond to `Conv` (basis conversion) and `ModDown` (modulus switching) in Appendix A of [Faster Bootstrapping for CKKS with Less Modulus Consumption](https://eprint.iacr.org/2025/1403.pdf). While investigating how `ModDown` could be implemented I also found the following: - The [LWE level basis conversion op](https://github.com/google/heir/blob/main/lib/Dialect/LWE/IR/LWEOps.td#L324) gets [lowered](https://github.com/google/heir/blob/main/lib/Dialect/LWE/Conversions/LWEToPolynomial/LWEToPolynomial.cpp#L631) to the [RNS level basis conversion](https://github.com/google/heir/blob/main/lib/Dialect/RNS/IR/RNSOps.td#L72). - The RNS level basis conversion op then gets lowered to further RNS and ModArith level operations [here](https://github.com/google/heir/blob/main/lib/Dialect/RNS/Transforms/LowerConvertBasis.cpp#L55) with Garner's algorithm. I propose the following fixes: - Remove the LWE level basis conversion op. The direct translation from LWE basis conversion to RNS basis conversion is an indication that basis conversion should live in the RNS dialect and not the LWE one. - Define a `ModUp` (basis extension from `Q` to `QP`) and a `ModDown` (modulus switching from `QP` to `Q`) operation on the LWE dialect following the standard naming convention from Section 2.2 of [HK19](https://eprint.iacr.org/2019/688.pdf). These ops will use the RNS level base conversion when lowered. - Replace the usages ([basis extension](https://github.com/google/heir/blob/main/lib/Dialect/CKKS/Transforms/DecomposeKeySwitch.cpp#L120) and [modulus switching](https://github.com/google/heir/blob/main/lib/Dialect/CKKS/Transforms/DecomposeKeySwitch.cpp#L146)) of the current LWE level basis conversion with the newly added `ModUp` and `ModDown` ops. - Reimplement the RNS level basis conversion lowering to use the fast base conversion algorithm ([Section 2.2](https://eprint.iacr.org/2019/688.pdf)), which is the standard algorithm used in CKKS (hybrid) key switching implementations. The Garner's algorithm implementation could optionally be kept as an exact basis conversion algorithm. If you think this is a good direction I can start putting together PRs for these. I am not confident about the compiler/MLIR side of things, so any suggestions are welcome. tagging @AlexanderViand, @crockeea, @j2kun
7 条评论