google / heir

A compiler for homomorphic encryption
https://heir.dev/
Apache License 2.0
354 stars 51 forks source link

LWEToPolynomial: Lowering RLWE encrypt and decrypt ops and Adding Randomness #762

Closed lawrencekhlim closed 3 months ago

lawrencekhlim commented 5 months ago

LWEToPolynomial: Lowering RLWE encrypt and decrypt ops and Adding Randomness

Background:

This issue discusses design decisions and changes for lowering LWE encrypt and decrypt to Polynomial. See

Remaining work:

Open questions:

For BGV:

Following the explanation of BGV in https://www.inferati.com/blog/fhe-schemes-bgv, here is the implementation we suggest for LWE to Polynomial.

a <- random polynomial sampled from uniform e <- random noise sampled from \Chi u <- random polynomial sampled from uniform

Plaintext message

In BGV, messages are stored in lower bits and error is scaled by the plaintext modulus t (to store error in the MSBs). Hence, the cleartext_start and cleartext_bitwidth are the same (which we arbitrarily wrote here as 10).

!pt = !lwe.rlwe_plaintext

Following that, we can calculate that the

cleartext_bitwidth = “log_2(t)” plaintext modulus


* Do we need the plaintext modulus t explicitly? (Instead of cleartext_bitwidth, we could have an attribute for plaintext_modulus). Note, this is not the implementation:

encoding = #lwe.polynomial_evaluation_encoding


### Public and Private Key Material

!rlwe_secret_key = !lwe.rlwe_secret_key !rlwe_public_key = !lwe.rlwe_public_key

In our implementation, lwe.rlwe_encrypt can take either a secret key or a public key, so lowering will decide whether we use the secret key encrypt operation or public key encrypt operation.

### New Dialect for Random sampling

#### We will need support for:

R_q: Uniform random over q \Chi: Discrete gaussian distribution (params: stdev, mean) R_2: Uniform ternary distribution {-1, 0, 1} U_2: Uniform binary distribution

* Generating sk (secret key) is usually from ternary distribution
* Generating pk (public key) uses both discrete gaussian distribution (\Chi) and uniform distribution R_q
* Secret Key Encryption uses both discrete gaussian distribution (\Chi) and uniform distribution R_q
* Public Key Encryption uses both discrete gaussian distribution (\Chi), uniform ternary distribution R_2
* Some of these may change, (for instance, the secret key distribution could change, based on attributes that we later add)

Two options for using randomness: (1) Create an RNG object with parameters (e.g. seeds, mean, stdev) from which we can generate samples or (2) No RNG object, simply use function calls to generate randomness, (3) Separating the seeding of an RNG object (the seed is not a parameter)

#### Create an RNG type:

%0 = random.discrete_gaussian_distribution { mean = ?, stdev = ?, prec = ?, seed = ? } : random.rng %1 = random.sample %0 : i32


#### No RNG object:

%0 = random.discrete_gaussian_distribution { mean = ?, stdev = ?, prec = ?, seed = ? } : %1 = random.sample %0 : i32


#### Seed:

%0 = random.discrete_gaussian_distribution { mean = ?, stdev = ?, prec = ?} : (optional): random.seed %0, %seed_value %1 = random.sample %0 : i32


### Centering Uniform Distributions: 

Sometimes coefficients mod q need to be centered at 0 and other times they can be in the range `[0, q)`. To represent this, we could either always choose to emit `[0, q)` and handle centering at zero with an explicit arithmetic subtraction, or add an attribute to the op (e.g. `signed`?).

#### Option 1: Generating R_2 with subtraction

%q = arith.constant q : i32 %0 = random.uniform_random_mod {modulus = q} : tensor<1024xi32> %cq = arith.floordivsi %q, %c2 : i32 %1 = arith.sub %0, %cq : tensor<1024xi32>


#### Option 2: Generating R_2 with a flag indicating signed

%0 = random.uniform_random_mod {modulus = q, signed} : tensor<1024xi32>


@j2kun: What does precision mean in SageMath’s discrete gaussian distribution? 
BGV

### Proposed implementation for encryption:
Following the algorithm described in [https://www.inferati.com/blog/fhe-schemes-bgv](https://www.inferati.com/blog/fhe-schemes-bgv), we propose this implementation where c1 and c2 are our ciphertext pair.

Calculate t = 2**(bitwidth) Branch on whether we have a public or private key: Public Key Encryption: %u = random.sample_ternary %e1 = random.sample_gaussian %e2 = random.sample_gaussian <- (Value getGaussianSample(rewriter)) %pk1 = tensor.extract %pk[0] %pk2 = tensor.extract %pk[1] %c1 = polynomial.add(polynomial.add(polynomial.mul (pk1, u), polynomial.mul(t, e_1)), M) %c2 = polynomial.add(polynomial.mul (pk2, u), polynomial.mul(t, e_2))) Private Key Encryption: Figure out later



@j2kun 
@AlexanderViand-Intel 
AlexanderViand-Intel commented 5 months ago

Thanks for the great writeup! There's a lot to go through here, so let me start with some high-level comments:

Following the explanation of BGV in https://www.inferati.com/blog/fhe-schemes-bgv [...]

I don't know what exact variant is presented there, but we should make sure to use the encryption approach from Revisiting Homomorphic Encryption Schemes for Finite Fields (ASIACRYT'21) which is a unified BFV/BGV analysis. I'm not sure if their BGV is also "improved", but I know their BFV encryption has lower noise growth than the original BFV formulation. This is also what OpenFHE does (See "§2.2.1, BFV Scheme" in the WAHC'22 OpenFHE paper) Though even if there is a difference to the lowering you have right now, it'll probably be a minor one.

Does the plaintext modulus always need to be a power of 2 (current RLWE encoding restricts to only powers of 2, which is not necessarily true because besides TFHE, the plaintext modulus might need to be prime)?

Indeed, the restriction should be removed as the plaintext (coefficient) modulus can be pretty much anything in RLWE schemes, depending on what encoding/packing approach is used.

Encoding Attributes

I'm having a bit of a hard time parsing this section. I wouldn't expect MSB/LSB (BGV/BFV) discussions to make a difference at the plaintext encoding stage yet, and I'm not quite sure what polynomial_evaluation_encoding is capturing, but that might just be me failing to get it.

Message/Cleartext -> Plaintext Encoding

I'd expect a plaintext_encoding attribute to either be more like a map, describing how the message, e.g., a tensor<Nxi...> {encoding = "mod t"} is converted into a rlwe_plaintext . Since it might be hard to do that generically, the attribute could also be more of an "enum" style indicator right now, with options such as coeff (for the trivial coefficient-wise packing) or crt/full/or some other name for the "default" SIMD packing we typically use in RLWE schemes, with the option to extend this with more advanced techniques later and/or add a map-like way to handle all kinds of encodings generically throughout the compiler.

Plaintext -> Ciphertext Encryption

This is really the stage at which I'd expect to see the MSB/LSB BGV/BFV difference to pop up, and where it also starts to be relevant to be able to track things such as the start/length of the plaintext bits and what happens to noise.

Supporting Generic RLWE Lowerings

Looking at this also prompted me to think about we need to preserve semantics across cleartext-plaintext-ciphertext, which is not directly related to this PR, but probably good to keep in mind as we design the key attributes that capture this.

If we think about what needs to happen to lower secret.generic(%t1, %t2 : ...) { ... arith.muli %st1, %st2 : tensor<8192xi32> ... }, my first realization is actually that we'd need to either homomorphically emulate the i32 overflow behavior (for a 100% "correct" lowering) or have a way to express what "deviation" between plaintext modulus t and $2^{32}$ the program can tolerate. (Cheap solution: pass parameters? More elegant might be to extend !secret.secret<type> with (optional) attributes that specify this, but then we'd also need to support that for the secretize slap-an-attribute-on-func-params "shortcut", as I'd guess most frontends will want to use something like that).

However, ignoring the t vs $2^{32}$ semantics issue right now, the next thing we need to know is the plaintext semantics. If we do a coeff encoding, then we cannot simply realize this as a multiplication between two plaintexts, but if we use standard SIMD packing, then this would be correct. Of course, even a plaintext-plaintext multiplication becomes a tensoring+relin when we want to do it homomorphically over encrypted values. In practice, the lowering actually goes straight from secret to bgv, but there's clearly some encoding<->lowering connection, as the current secret-to-bgv lowering would produce incorrect code if used with coeff packing.

Randomness

I'll need to go back and re-read your part on randomnessmore carefully, but here's a few initial things I noticed:

Randomness Seeding

(3) Separating the seeding of an RNG object (the seed is not a parameter)

We generally want to be able to both choose and explicitly encode into the IR the PRNG and seed used for "public" randomness (e.g., a). This is a requirement for the common ciphertext compression (basically, instead of storing/sending a, you just agree on an PRNG and send the seed around) and since different targets will support different PRNGS (or none at all), this needs to be something we can parameterize.

Modular Representation

It's interesting that the [0, q) vs [-q/2, q/2) representation issue pops up again here. We discussed this in the past in the context of lowering polynomial to LLVM, thinking about modular arithmetic and (lack of) semantics for arith and unsigned integers. This is probably worth its own discussion/issue, as it's something we need to settle ASAP before we start having different parts of the compiler assume different things. One issue I could see is if HEIR assumes a certain representation and the target (library, arith->LLVM lowering, HW accelerator) has a different reprsentation, any metadata/constants that are in the IR will be "wrong".

j2kun commented 5 months ago

I haven't had a chance to read through everything, but on the MSB/LSB question, I consider that part of plaintext encoding, not encryption. But I imagine the references combine the two into one step. For comparison, the equivalent operation is part of LWE encoding in CGGI, so it might be worth keeping them aligned for consistency.

On Tue, Jul 2, 2024, 2:36 AM Alexander Viand @.***> wrote:

Thanks for the great writeup! There's a lot to go through here, so let me start with some high-level comments:

Following the explanation of BGV in https://www.inferati.com/blog/fhe-schemes-bgv [...]

I don't know what exact variant is presented there, but we should make sure to use the encryption approach from Revisiting Homomorphic Encryption Schemes for Finite Fields (ASIACRYT'21) https://eprint.iacr.org/2021/204 which has lower noise growth than the original BFV formulation. This is also what OpenFHE does (See "§2.2.1, BFV Scheme" in the WAHC'22 OpenFHE paper https://dl.acm.org/doi/pdf/10.1145/3560827.3563379) Though if there is a difference to the lowering you have right now, it'll probably be a minor one.

Does the plaintext modulus always need to be a power of 2 (current RLWE encoding restricts to only powers of 2, which is not necessarily true because besides TFHE, the plaintext modulus might need to be prime)?

Indeed, the restriction should be removed as the plaintext (coefficient) modulus can be pretty much anything in RLWE schemes, depending on what encoding/packing approach is used. Encoding Attributes

I'm having a bit of a hard time parsing this section. I wouldn't expect MSB/LSB (BGV/BFV) discussions to make a difference at the plaintext encoding stage yet, and I'm not quite sure what polynomial_evaluation_encoding is capturing, but that might just be me failing to get it. Message/Cleartext -> Plaintext Encoding

I'd expect a plaintext_encoding attribute to either be more like a map, describing how the message, e.g., a tensor {encoding = "mod t"} is converted into a rlwe_plaintext . Since it might be hard to do that generically, the attribute could also be more of an "enum" style indicator right now, with options such as coeff (for the trivial coefficient-wise packing) or crt/full/or some other name for the "default" SIMD packing we typically use in RLWE schemes, with the option to extend this with more advanced techniques later and/or add a map-like way to handle all kinds of encodings generically throughout the compiler. Plaintext -> Ciphertext Encryption

This is really the stage at which I'd expect to see the MSB/LSB BGV/BFV difference to pop up, and where it also starts to be relevant to be able to track things such as the start/length of the plaintext bits and what happens to noise. Supporting Generic RLWE Lowerings

Looking at this also prompted me to think about we need to preserve semantics across cleartext-plaintext-ciphertext, which is not directly related to this PR, but probably good to keep in mind as we design the key attributes that capture this.

If we think about what needs to happen to lower secret.generic(%t1, %t2 : ...) { ... arith.muli %st1, %st2 : tensor<8192xi32> ... }, my first realization is actually that we'd need to either homomorphically emulate the i32 overflow behavior (for a 100% "correct" lowering) or have a way to express what "deviation" between plaintext modulus t and $2^{32}$ the program can tolerate. (Cheap solution: pass parameters? More elegant might be to extend !secret.secret with (optional) attributes that specify this, but then we'd also need to support that for the secretize slap-an-attribute-on-func-params "shortcut", as I'd guess most frontends will want to use something like that).

However, ignoring the t vs $2^{32}$ semantics issue right now, the next thing we need to know is the plaintext semantics. If we do a coeff encoding, then we cannot simply realize this as a multiplication between two plaintexts, but if we use standard SIMD packing, then this would be correct. Of course, even a plaintext-plaintext multiplication becomes a tensoring+relin when we want to do it homomorphically over encrypted values. In practice, the lowering actually goes straight from secret to bgv, but there's clearly some encoding<->lowering connection, as the current secret-to-bgv lowering would produce incorrect code if used with coeff packing. Randomness

I'll need to go back and re-read your part on randomization more carefully, but here's a few initial things I noticed: Randomness Seeding

(3) Separating the seeding of an RNG object (the seed is not a parameter)

We generally want to be able to both choose and explicitly encode into the IR the PRNG and seed used for "public" randomness (e.g., a). This is a requirement for the common ciphertext compression (basically, instead of storing/sending a, you just agree on an PRNG and send the seed around) and since different targets will support different PRNGS (or none at all), this needs to be something we can parameterize. Modular Representation

It's interesting that the [0, q) vs [-q/2, q/2) representation issue pops up again here. We discussed this in the past in the context of lowering polynomial to LLVM, thinking about modular arithmetic and (lack of) semantics for arith and unsigned integers. This is probably worth its own discussion/issue, as it's something we need to settle ASAP before we start having different parts of the compiler assume different things. One issue I could see is if HEIR assumes a certain representation and the target (library, arith->LLVM lowering, HW accelerator) has a different reprsentation, any metadata/constants that are in the IR will be "wrong".

— Reply to this email directly, view it on GitHub https://github.com/google/heir/issues/762#issuecomment-2202564314, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAS2PKUBQ4PB7MFRZINHOTDZKJYB5AVCNFSM6AAAAABKGHJ4D6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEMBSGU3DIMZRGQ . You are receiving this because you were mentioned.Message ID: @.***>

asraa commented 5 months ago

I haven't had a chance to read through everything, but on the MSB/LSB question, I consider that part of plaintext encoding, not encryption. But I imagine the references combine the two into one step.

Yep! That's correct, FWIW this issue was just recapping the encoding type.

AlexanderViand-Intel commented 5 months ago

I'm not super familiar with the CGGI naming conventions, but "LWE Encoding" sounds like it would already be operating at the ciphertext level ("with errors"), despite the naming?

Maybe we're talking about something else, but I was thinking about the BFV-vs-BGV MSB/LSB difference, i.e., whether you map from the plaintext space $R_t$ (plaintext poly with coefficient mod $t$) to the ciphertext space $R_q^2$ by either (a) as in BFV, multiplying your plaintext with a large scalar (i.e., $\Delta = \lfloor \frac{q}{t} \rfloor$) to "put it in the MSB" or (b) multiplying the noise by $t$, effectively putting the plaintext in the LSB. Oh, and of course (c) let plaintext and noise overlap in the LSB, as in CKKS.

Since this only happens during encryption, it seems a bit odd to have this encoding stored on the plaintext already. For example, we could take the same Plaintext and encrypt it with BFV, BGV and CKKS, each of which would apply different shifts/scales. Even within the same system, if we encrypted the same Plaintext at different levels (i.e., to avoid later mod-switches), the offsets/scales would be different for each encryption.

Btw, the inferati blog uses "Plaintext Encoding" to refer to coeff-packing vs SIMD-packing, and that usage seems to align more with how I see that term used. (I.e., for things such as lwe.polynomial_evaluation_encoding and lwe.polynomial_coefficient_encoding)

j2kun commented 5 months ago

"Since this only happens during encryption"

My feeling is that this can happen during encoding as well, but I suppose in the RLWE schemes what differs is that the scaling depends on q which is not part of the plaintext space, so it must happen during encryption.

lawrencekhlim commented 5 months ago

Here are a few things we (meaning @asraa and @lawrencekhlim me) decided or discussed:

Randomness

We decided on our Randomness dialect, based on the fact that seeding may need to be determined during runtime and examining other libraries' implementations of randomness:

So, in summary, our randomness might look like the following:

// random.prng (type) provides psuedorandom numbers
%0 = random.init_prng %seed_value : !random.prng

// random.distribution (type) provides values in the given distribution 
%1 = random.discrete_gaussian_distribution { stdev = ?, mean = ?, etc} %0 : prng -> distribution
%1 = random.uniform_discrete_distribution { range = [-1, 2] } %0 : prng -> distribution

// random.sample : random.distribution -> AnyInteger
%2 = random.sample %1 : !random.prng -> int32

Encoding Schemes

We actually already consider the different encoding/decoding schemes. Inside lib/Dialect/LWE/IR/LWEAttributes.td we have different attributes for different encoding schemes. For instance RLWE_PolynomialCoefficientEncoding represents RLWE plaintexts with coefficient encodings, RLWE_PolynomialEvaluationEncoding represents RLWE plaintexts with evaluations at fixed points, and so on with RLWE_InverseCanonicalEmbeddingEncoding.

More Efficient Encryption Approach

Ack, regarding using the encryption approach from Revisiting Homomorphic Encryption Schemes for Finite Fields (ASIACRYT'21). I'll have to take some time to read through it though.

j2kun commented 5 months ago

potential conflicts with the ternary uniform random [-q/2, q/2) and mod q representation of [0, q)

What is the conflict here? The operations on the underlying bits are identical regardless of the choice of representation, no? That's the reason why a 32-bit adder circuit is identical for unsigned and (2s-complement) signed i32s. I feel as though we should be able to choose the canonical representative relatively late in the pipeline.

lawrencekhlim commented 5 months ago

"Conflict" isn't the best word so much as needing to support both ternary and mod q requires us to make design decisions on how to do so.

lawrencekhlim commented 3 months ago

Closing issue since we've addressed most everything here, plus or minus several details about encoding or randomness that we may want to adjust.