google
lwe-to-polynomial: migrate encrypt/decrypt lowering to support RNS
PR #1190 migrated most of the types to NewLWEType, however, for lwe-to-polynomial, there are some non-trivial changes that needs to be done.
This also relates to #882
NewLWEPlaintext now has a PlaintextSpace.Ring (for example Z_65537 / (X^N + 1)), which is different from NewCiphertext's CiphertextSpace.Ring (typically RNS)
- lowering should convert plaintext ring to rns ring
- lowering should properly handle modulus switching from
Z_65537toZ_qi(ciphertext rns modulus), which needsmod_arithdialect support, and it is non-trivial. - lowering should handle scaling factor for CKKS scheme.
This issue has 5 outstanding TODOs:
- lib/Dialect/LWE/Conversions/LWEToPolynomial/LWEToPolynomial.cpp:56: properly lower LWEType (often RNS) to PolynomialType.
- lib/Dialect/Polynomial/Conversions/PolynomialToModArith/PolynomialToModArith.cpp:774: support RNS lowering
- lib/Dialect/Polynomial/Conversions/PolynomialToModArith/PolynomialToModArith.cpp:843: support RNS lowering
- tests/Dialect/BGV/Conversions/bgv_to_polynomial/BUILD:9: support RNS lowering
- tests/Dialect/Polynomial/Conversions/heir_polynomial_to_llvm/BUILD:9: support RNS lowering
This comment was autogenerated by todo-backlinks
![github-actions[bot] avatar](https://avatars.githubusercontent.com/in/15368?v=4 "Published by a pub.dev verified publisher"){kind=small}
IIUC @asraa is working on this, and note it's a blocker for https://github.com/google/heir/pull/2013/files
The remaining part of this (that https://github.com/google/heir/pull/2036 pushes) is that RNS types are not supported in polynomial-to-mod-arith now. Since poly coeff types can be RNS, poly operations are still operating on the RNS polynomial types. A further pass, e.g. polynomial-to-mod-arith should decompose the RNS polynomials into the individual limbs.
A further pass, e.g. polynomial-to-mod-arith should decompose the RNS polynomials into the individual limbs.
It seems reasonable to define mod_arith ops to work on rns types as if they were vector types... Then maybe we could get away with inserting an "rns-unroll" type pass before mod-arith-to-arith. Thoughts?
A further pass, e.g. polynomial-to-mod-arith should decompose the RNS polynomials into the individual limbs.
It seems reasonable to define mod_arith ops to work on rns types as if they were vector types... Then maybe we could get away with inserting an "rns-unroll" type pass before mod-arith-to-arith. Thoughts?
Yep, something like LimbwiseMappable would be great :)
Oooh I like this!
I'm just curious - does hardware natively support a vector of different modular operations? Or do they break them down into individual single mod Z_p operations?
I'm just curious - does hardware natively support a vector of different modular operations? Or do they break them down into individual single mod Z_p operations?
This might differ from vendor to vendor, but I think for HERACLES, we eventually want to re-order things to be grouped by RNS limb, but I think LimbwiseMappable would actually make that easy - instead of lowering each op to a loop over the limbs, we could just take a chunk of LimbwiseMappable operations and lower them to a combined loop over the limbs.
It seems reasonable to define mod_arith ops to work on rns types as if they were vector types... Then maybe we could get away with inserting an "rns-unroll" type pass before mod-arith-to-arith. Thoughts?
mod_arith ops (add/sub/mul) already support rns types (#1744), though it does not have a lowering for it; for mod_arith mod_switch op, it supports rns lowering, though not in a rns-unroll style (#1928)
Ah right, yes. mod_arith ops take RNS types too. I think then maybe an RNS-unroll transform pass at the mod_arith level is what this should be