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!

[ZenithalHourlyRate]

I was trying to reuse this PR and found either the type way !polynomial.rns or the attr way RNSRingAttr for !polynomial.polynomial not satisfying.

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

I like the idea of supporting RNS also for integers/ctxts. I think the current NewLWECiphertextType still can not handle RNS quite well. And the current ModArith dialect should also handle rns integer arithmetic.

I think this can be technically done in the following way:

• Introduce ModularIntegerType for ModArith, and RNSIntegerType is naturally tuple<ModularIntegerType>
• Allow polynomial.ring<coefficientType=(i32|f32|Zp|tuple<Zp1, Zp2>)> so coefficientModulus is no longer needed for polynomial type.
• NewLWECiphertextType uses such ring.

[j2kun]

FYI, I'm planning to start working on this actively next week when I'm back from leave.

[AlexanderViand-Intel]

I think this can be technically done in the following way:

• Introduce ModularIntegerType for ModArith, and RNSIntegerType is naturally tuple<ModularIntegerType>
• Allow polynomial.ring<coefficientType=(i32|f32|Zp|tuple<Zp1, Zp2>)> so coefficientModulus is no longer needed for polynomial type.
• NewLWECiphertextType uses such ring.

This (minus the modarith, which didn't exist back then) is actually how I implemented RNS for my first hacky proof-of-concept. It has the advantage of being really easy to implement, and enabling very easy lowerings. However, I don't think it's a good solution overall, as it (a) makes the polynomial type messy, breaking the abstraction in a really weird FHE-specific way and (b) doesn't really allow us to re-use any of the RNS handling logic across polys and integers (as far as I can see)

EDIT: I also realized there's something a bit subtle going on with the polynomial type that maybe our current documentation doesn't make clear. !polynomial.polynomial's coeffiecientType indicates the intended storage type for the datastructure representing the polynomial, which is distinct from the "logical" type (which we could just derive from the coefficientModulus anyway). I think there's a good reason for that, but I actually don't remember right now what that is 😅

[AlexanderViand-Intel]

Closing in favor of #1126 and related changes.