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

optimize-relinearization: Remove requirement that arguments have equal key-basis degrees #1295

Merged
merged 1 commit into from
Jan 27, 2025

Conversation

j2kun
Copy link
Collaborator

@j2kun j2kun commented Jan 23, 2025

This PR modifies the ILP for optimize-relinearization to permit the possibility that the operands to an op like arith.addi have mixed key-basis degrees. E.g., you can add a ciphertext with basis (1, s) to another with basis (1, s, s^2).

Based on the discussion in #1284, Lattigo and OpenFHE both support this for BGV.

To hedge against the possibility a future backend does not support this, I made the relaxation optional via a new pass option (allow-mixed-level-ops).

The impact of this relaxation on the ILP is not as trivial as I had hoped. When removing the constraints on key basis degree for operands, I had to add new constraints that the output degree is at least as big as each input degree. Ideally this comes paired with an objective term that makes the output degree as small as possible, which would implement a result_degree = max(arg0_degree, ..., argN_degree) function. However, adding this to the objective causes the solver to "cheat" by inserting relinearization terms earlier than necessary to as to reduce cost of the new terms.

The "right" solution here is to properly cost model the solver's choices (#1018). Instead of minimizing the total number of relin ops, have a cost for each op at each level, as well as the cost of doing a relin op from each starting level.

In the interim, we simply ignore the fact that the solver may assign a larger key-basis-degree for the output of an add op than necessary. It doesn't seem to affect the correctness of any results; hallucinating larger key-basis degree doesn't change whether or not you need to relinearize, just the mechanics of how to relinearize, which is handled opaquely by the backend anyway. In particular, if the output degree of an op were truly 1, but the solver decided to make it 2, this would require inserting a relin op when it's not necessary, so the solver naturally avoids this situation. But if the degree were 2 and the solver hallucinated it should be 3, we still relin.

Fixes #1284

@j2kun j2kun changed the title Reproducing test for #1284 optimize-relinearization: Remove requirement that arguments have equal key-basis sizes Jan 24, 2025
@j2kun j2kun changed the title optimize-relinearization: Remove requirement that arguments have equal key-basis sizes optimize-relinearization: Remove requirement that arguments have equal key-basis degrees Jan 24, 2025
Copy link
Collaborator

@ZenithalHourlyRate ZenithalHourlyRate left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also remember to update the relinearization_ilp.md.

@j2kun
Copy link
Collaborator Author

j2kun commented Jan 27, 2025

Updated the doc as well

Copy link
Collaborator

@ZenithalHourlyRate ZenithalHourlyRate left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For relinearization_ilp.md there is one more place mentioning the same degree requirement, i.e.

Addition can operate in any key basis (though all inputs must have the same key basis),

@AlexanderViand-Intel
Copy link
Collaborator

This PR modifies the ILP for optimize-relinearization to permit the possibility that the operands to an op like arith.addi have mixed key-basis degrees. E.g., you can add a ciphertext with basis (1, s) to another with basis (1, s, s^2).

Based on the discussion in #1284, Lattigo and OpenFHE both support this for BGV.

To hedge against the possibility a future backend does not support this, I made the relaxation optional via a new pass option (allow-mixed-level-ops).

I think the main issue here is for CKKS and the GHS variant of BGV, where the messages are scaled. In those settings, we can't just do $(c_0, c_1) + (c'_0, c'_1, c'_2) = (c_0 + c'_0, c_1 + c'_1, c'_2)$ (as we can, afaik, in traditional BGV and in BFV), because the message in ctxt c will be scaled by a different factor than that in ctxt c'. While being able to toggle mixed operations off is good, I think it's not actually sufficient to prevent similar scaling-related issues all the time.

So maybe we don't even need the toggle, but just a solution to #1169 (tracking the message scale) + making the modswitch mgmt aware of the need to adjust the scales?

@j2kun
Copy link
Collaborator Author

j2kun commented Jan 27, 2025

Squashed and marking pull_ready

@j2kun j2kun added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Jan 27, 2025
@copybara-service copybara-service bot merged commit b9500b6 into google:main Jan 27, 2025
8 checks passed
@j2kun j2kun deleted the issue-1284 branch February 13, 2025 05:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

operation-rebalancer pass modifies IR so as to require additional relinearization ops
4 participants