Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Lattigo: Correctly lowering 1-dim bgv.rotate to RotateColumns/RotateRows for full packing #1186

Open
ZenithalHourlyRate opened this issue Dec 13, 2024 · 8 comments
Labels
dialect: lattigo Issues concerning the Lattigo dialect

Comments

@ZenithalHourlyRate
Copy link
Collaborator

Lattigo natively exposes the Z_t/(X^N+1) plaintext space as a SIMD packing of a matrix of size N/2 * 2, which element in Z_t, which is due to the nature of automorphism.

The Galois automorphisms form a group (under composition), which is isomorphic to Zn/2 × Z2. The first factor is generated by gal_elt = 3, and the second factor is generated by gal_elt = m − 1

https://www.microsoft.com/en-us/research/uploads/prod/2017/11/sealmanual-2-3-1.pdf , section 5.6

Their API allows RotateColumns by k (either to the left or right, size N/2) or RotateRows (swapping rows).

In the meantime, OpenFHE exposes it just as a vector of N, and provides an API generating indexes of rotation of size N.

If we just pack half of the plaintext, it is OK to just use RotateColumns. But if we do full packing for BGV for Lattigo, we need to correctly combine the Lattigo API to provide a 1-dim view of the data.

@ZenithalHourlyRate
Copy link
Collaborator Author

OpenFHE also shows a n/2 by 2 feature, and when rotate right by 1, these two rows does not mix, so we does not have a generic 1-dim bgv.rotate for index > n / 2.

For full packing, tensor_ext.rotate has distinct semantic as bgv.rotate. The workaround is to always half-full-packing for BGV for now. (Note that in cyclic repetition code I did a half packing)

For example, the following code get the full batching for BGV, but when rotated right by 1, we get

auto batchSize = cryptoContext->GetCryptoParameters()->GetEncodingParams()->GetBatchSize();
    std::vector<int64_t> vectorFull;
    for (size_t i = 0; i < batchSize; i++) {
        vectorFull.push_back(i);
    }
    Plaintext pt = cryptoContext->MakePackedPlaintext(vectorFull);
    auto ct      = cryptoContext->Encrypt(keyPair.publicKey, pt);
    auto ct_rot  = cryptoContext->EvalRotate(ct, -1);
    Plaintext dec;
    cryptoContext->Decrypt(keyPair.secretKey, ct_rot, &dec);
    std::cout << "[";
    for (size_t i = 0; i < 3; i++) {
        std::cout << dec->GetPackedValue()[i] << " ";
    }
    std::cout << "... ";
    for (size_t i = batchSize / 2 - 3; i < batchSize / 2; i++) {
        std::cout << dec->GetPackedValue()[i] << " ";
    }
    std::cout << "]\n";
    std::cout << "[";
    for (size_t i = 0; i < 3; i++) {
        std::cout << dec->GetPackedValue()[i + batchSize / 2] << " ";
    }
    std::cout << "... ";
    for (size_t i = batchSize / 2 - 3; i < batchSize / 2; i++) {
        std::cout << dec->GetPackedValue()[i + batchSize / 2] << " ";
    }
    std::cout << "]";
Batch size: 16384
[8191 0 1 ... 8188 8189 8190 ]
[16383 8192 8193 ... 16380 16381 16382 ]

instead of

[16383 0 1 ... 8188 8189 8190 ]
[8191 8192 8193 ... 16380 16381 16382 ]                                                                                             

Openfhe do allow a rotation amount larger than n/2, but it is achieved by first rotate rows then rotate corresponding columns to get the corresponding element in place.

See also on how to compute automorphism index for index > n/2

https://github.com/openfheorg/openfhe-development/blob/7b8346f4eac27121543e36c17237b919e03ec058/src/core/lib/math/nbtheory2.cpp#L212-L225

@j2kun
Copy link
Collaborator

j2kun commented Dec 14, 2024

Just briefly browsing the code, it appears that the Lattigo API hence requires you to apply at least two distinct automorphisms to achieve a rotation of bigger than N/2. Is OpenFHE doing the same, but hiding it behind the API? Or is OpenFHE somehow precomputing the composite automorphism element ahead of time and then requiring only a single application?

I suspect this will have a material performance impact if there's a difference.

@ZenithalHourlyRate
Copy link
Collaborator Author

ZenithalHourlyRate commented Dec 14, 2024

it appears that the Lattigo API hence requires you to apply at least two distinct automorphisms to achieve a rotation of bigger than N/2.

After deeper inspection, it appears that if we manually call automorphism we can get a "rotation" bigger than N/2

// logN = 14
// galElement = M - 5, means swapping row then rotating left by 1, or just "rotating left" by N/2 + 1.
v14 := v10.GenGaloisKeyNew(32768-5, v11)
// input with index
v19 := make([]int64, 16384)
for i := range v19 {
  v19[i] = int64(i)
}

err := v0.Automorphism(v1, 32768-5, v2)

fmt.Println("First three elements:", v30[:3])
fmt.Println("Middle six elements:", v30[len(v30)/2-3:len(v30)/2+3])
fmt.Println("Last three elements:", v30[len(v30)-3:])

We get

First three elements: [8193 8194 8195]
Middle six elements: [16382 16383 8192 1 2 3]
Last three elements: [8190 8191 0]

Or is OpenFHE somehow precomputing the composite automorphism element ahead of time and then requiring only a single application?

OpenFHE calls only one automorphism with precomputed composite automorphism element

@j2kun
Copy link
Collaborator

j2kun commented Dec 14, 2024

First three elements: [8193 8194 8195]
Middle six elements: [16382 16383 8192 1 2 3]
Last three elements: [8190 8191 0]

That example isn't quite rotating left by N/2 + 1 though, the 8192 and 0 are out of place (because they are the rotations of each row, separately). An arbitrary rotation by k is not expressible in Z/(N/2)Z x Z/2Z. (For example, it's easier to see that Z/2Z x Z/2Z cannot represent a rotation by 3)

I am embarrassed to admit I did not know about this automorphism group structure in FHE before now (I had not bothered to learn this part of the literature deeply enough).

So if we want an arbitrary rotation by k, we will need to decompose it into a sequence of rotations and masks. E.g., to rotate C by N/2 + 1 (let N = 8 for this example):

Swap the rows + rotate by 1

5 6 7 4
1 2 3 0

Mask the first N/2 - 1 elements in each row (5-7 and 1-3) and zero out the rest; call this C1

5 6 7 0
1 2 3 0

Rotate C (original ciphertext) by 1

1 2 3 0
5 6 7 4

Mask the last column, call this C2

0 0 0 0
0 0 0 4

Add C1 + C2 to get the final result.

This construction gives a generic rotation with 2 galois automorphisms + 2 ciphertext-plaintext muls + 1 ciphertext addition. It's a little bit disappointing, but perhaps this also suggests we should define rewrite patterns at the level of galois automorphism (e.g., composing automorphisms by finding the relevant composite index) and hoisting (the same ciphertext with many automorphisms applied).

Perhaps this suggests we should have a proper galois automorphism op that tensor_ext.rotate lowers to? Though knowing about the N/2 x 2 representation of the ciphertext is a whole other issue with the type system...

@ZenithalHourlyRate
Copy link
Collaborator Author

have a proper galois automorphism op that tensor_ext.rotate lowers to

I think instead this is highly related to how the "packing" is done. tensor_ext.rotate is intended to be operating on 1-dim tensor, yet if full packing bgv ciphertext can not be expressed as a 1-dim tensor. Instead of the composite process above, we should at the first place not lower a large 1-dim tensor to full packing BGV ciphertext. Or we can view rows as 2 independent tensor and never allow interaction between them, for the earlier "packing" pipeline.

@j2kun
Copy link
Collaborator

j2kun commented Dec 14, 2024

Perhaps this is something we should measure, but I suspect never using a full packing would be too harmful enough to performance to do forever.

In the mean time, to limit the scope of this issue sufficiently, should we just implement the two-step method as a rewrite pattern at the BGV level? Or during the lowering from BGV-to-API? Maybe we could add a new bgv.rotate_rows and bgv.rotate_columns ops alongside bgv.rotate, and they would just interpret a degree-N ciphertext as having a 2 x N/2 layout without encoding it specifically in the type (or maybe specify it with a type attribute? though that then requires a dialect conversion pass for the type converter).

@ZenithalHourlyRate
Copy link
Collaborator Author

Maybe we could add a new bgv.rotate_rows and bgv.rotate_columns ops alongside bgv.rotate, and they would just interpret a degree-N ciphertext as having a 2 x N/2

I prefer to wait until we have a complete "packing" pipeline which can correctly identify sparsely packing/half-packing/full-packing, then we can have corresponding lowering. I particularly does not like the current --mlir-to-bgv pipeline parameter ciphertext-degree/poly-mod-degree as it is just slot number instead of true ciphertext degree.

Copy link

This issue has 1 outstanding TODOs:

This comment was autogenerated by todo-backlinks

@ZenithalHourlyRate ZenithalHourlyRate added the dialect: lattigo Issues concerning the Lattigo dialect label Dec 17, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
dialect: lattigo Issues concerning the Lattigo dialect
Projects
None yet
Development

No branches or pull requests

2 participants