RNS Support

[AlexanderViand-Intel]

This a DRAFT to elicit design feedback on how we should handle RNS support in HEIR. Towards that goal, it simply switches most things to RNS (rather than supporting RNS and non-RNS), and it also removes various verifiers/etc for ease of implementation.

Types

This uses an explicit RNS type, which is essentially a "Tuple" of Types. Note that the builtin Tuple type is not suitable, as it is not a valid tensor element type, and we want tensors or RNS-polynomials as they are the natural lowering for an (RNS) RLWECiphertext).

Rather than adding !rns.rns<...> and our RNS dialect to the llvm repo, this adds an rns type to the polynomial dialect (See the changes to LLVM in my fork here, but see also below for ideas how to avoid that.

In order to support polynomial ops on (a) normal polynomials (b) RNS polynomials and (c) tensors of both, it also adds a new TypeConstraint and uses it for Polynomial operations.

In order to track that BGV/LWE operations are happening in RNS (which changes things such as how a "modswitch" is lowered rather dramatically), it currently introduces an RNSRLWECiphertext. This is not very elegant, but I'm not quite sure how to make it nicer. Maybe the "RNS-ness" can instead become an Attribute on RLWECiphertext? See also #876 (PR on LWE Attrs).

Passes

This PR adds an --expand-rns pass that is similar to --convert-elementwise-to-affine (which can be used to convert polynomial ops on tensors of (rns-of) polys to loops with operations on "scalar" (rns-of) polynomials). It also updates LWE-to-polynomial to support RNS polynomials.

The intended order of passes is -<scheme>-to-LWE --LWE-to-polynomial (which will create poly operations on tensor<rns<poly1,poly2,poly3>>) followed by --convert-elementwise-to-affine , possibly --convert-tensor-to-scalars (from PR #763) and finally --expand-rns.

Finally, the PR switches the --secret-to-bgv pass to RNS, and changes the pass options from asking for the number of bits in the ctxt modulus to asking for a list of moduli.

[j2kun]

🤩

Just FYI, I have commit access in LLVM, so I can review and submit any upstream changes you have to polynomial dialect

[AlexanderViand-Intel]

I've cleaned up the diff and the commit history a bit (should be free of unrelated changes now) + updated the PR description above.

Below, I'll try and summarize some of the insights from the discussion at today's meeting:

• While we experiment with different designs, we can use the same "bazel patch trick" we used for the build hotfixes to let us try out changes to polynomial (on HEIR main ) before committing to upstreaming them. Since there's no-one else making changes to the Polynomial dialect upstream at the moment, the patch should not or only rarely fail to apply.

• Since upstreaming a dialect like RNS with a single type and maybe two ops (decompose/recompose) might not be easy, but having a !polynomial.rns type is also bad, as we want to support RNS on integers/ctxts/etc, too, we realized we might be able to switch polynomial operations to require a certain interface on its operand's types rather than specific types. We could then maybe even have interface methods that remove the need for passes like --expand-rns by letting the polynomial op lowerings hook a callback into some kind of "rewrite_for_each_element" thing. If we can augment tensor with these interface methods, too (assuming one can dynamically add interfaces to builtin stuff?), we could maybe even remove the need for --convert-elementwise-to-affine.

• Partially related: for representing large parameter sets (e.g., 1000+ bits of ctxt modulus) in textual IR, we might want to have custom printer/parsers for the modulus/moduli that use hex/base64 to make it less verbose.

• RNS already makes a difference for high-level things such as noise analysis, though rather than making whether we'll lower to an RNS version of BGV or a non-RNS version of BGV from the BGV-level IR, we could simply have RNS/non-RNS analysis and RNS/non-RNS lowerings (or, equivalently, options on unified analyses/passes)

Please let me know if I missed something!

[j2kun]

I think a good approach to avoiding duplication here would be to have RLWEParams take a more generic attribute for the ring. Instead of a ring attr, it could be a generic Attribute with a verifier, or AnyAttributeOf constraint for the legal RNS/non-RNS ring attirbute options.

[AlexanderViand-Intel]

Referencing #876 as I guess this should be taken into consideration in the LWE Attr&Types overhaul :)