google
Adding CKKS to HEIR
Adding CKKS to HEIR
Background
This issue describes design decisions and changes for creating a CKKS Dialect. CKKS is a FHE scheme that enables homomorphic floating point operations, and is especially useful for machine learning applications where the absolute value of the floating point value matters less than the approximate value. In this issue, we will not discuss how to lower CKKS to Polynomial or optimizations or bootstrapping, scheme switching, and other miscellaneous operators, but focus on OpenFHE's API for CKKS and how we might consider developing a CKKS dialect.
Needed work
- [x] Implement an initial CKKS Dialect and Ops. (See [lib/Dialect/BGV/IR])
- #502
- #835
- See https://github.com/google/heir/tree/main/lib/Dialect/BGV/IR) for an example.
- [ ] Support the Scale Modulus Size parameter (possibly adjustments to Secret?)
- [x] #501
- [x] Add lowerings for Secret to CKKS.
- See SecretToBGV for an example.
- [x] Add lowerings for CKKS to OpenFHE.
- See BGVToOpenfhe for an example.
- [x] Add lowerings for Secret to CKKS.
- [ ] Add CKKS bootstrapping operations
- [ ] Add lowerings for CKKS to Polynomial
Open Questions
- [x] How should we support the Scale Modulus Size parameter (needed for CKKS) in HEIR?
OpenFHE API
Here, we examine the OpenFHE API to see how best we can create a CKKS Dialect and what we need to support. Here are several examples of using CKKS in OpenFHE (some with C++ and some with Python):
- simple-real-numbers.cpp
- simple-real-numbers.py
- advanced-ckks-bootstrapping.cpp
- simple-ckks-bootstrapping.py
Parameters
To use CKKS in OpenFHE, one first creates a CryptoContext by specifying its parameters. In simple-real-numbers.cpp and simple-real-numbers.py, those parameters are
- Multiplicative Depth
- Scale Modulus Size
- Batch Size
- Security level (how many bits of security?)
| Parameter | Description |
|---|---|
| Multiplicative Depth | Indicates how many layers of homomorphic multiplication are desired |
| Scale Modulus Size | Indicates how much cleartext values are multiplied by before encoding. This scaling factor affects how many bits of precision (after the floating point) are available. A scale modulus Size of 50 means multiplying by 2^50. |
| Batch Size | Indicates how many cleartext values are packed into a plaintext / ciphertext. Note that for a ring dimension of N, the maximum batch size for CKKS is N/2 (because of how CKKS works) |
| Security Level | Indicates how many bits of security are desired (i.e. 128 bits of security) |
Note that we already have implementations to solve for the multiplicative depth (see lib/Analysis/MulDepthAnalysis, which will be useful for lowering CKKS).
It is still an open question as to how we will support these parameters right now. It's also possible that our interface to OpenFHE will not be these four parameters, but many others. In other words, those (meta-)parameters decide various other parameters such as plaintext modulus, ciphertext modulus, polynomial size, etc that we may choose based on the available hardware optimizations and software operations that will be run.
Operations
Again, following simple-real-numbers.cpp and simple-real-numbers.py, some available operations are
- Encoding (a packed cleartext to plaintext)
- Decoding (a packed plaintext to cleartext)
- Encryption (of an plaintext to ciphertext)
- Addition (of two ciphertexts)
- Ciphertext-Plaintext Addition
- Subtraction (of two ciphertexts)
- Multiplication (of two ciphertexts) (and Relinearize)
- Scalar Multiply (of a ciphertext and a floating point)
- Rotation (of a ciphertext)
- Negation
Given the obvious similarity of operations between the CKKS operations and BGV operations, the CKKS dialect and Ops can be made in the same way that BGV's dialect is constructed here (BGVOps.td).
For details on how to encode and decode messages, see this blog post. Note that this is where the scale modulus size parameter is needed (delta in the blog post).
WIP summary of what needs to be done. Let me know if I am missing anything. I will update this issue over time.
Scale and batch size will be needed for sure, but I think batch size will probably already be present in the IR attributes from how we handle packing.
The mul depth is already handled by Jaeho's analysis and the associated OpenFHE client configuration pass, so we can port that or generalize it to CKKS.
The security level is the same problem with the other OpenFHE lowerings, that this parameter causes OpenFHE to set lower level parameters that we'd ideally (long term) have more direct control over. So we can probably just set it to 128 bits and punt on that for the scope of supporting basic CKKS programs
I'd like to talk at the HEIR WG meeting this week on the LWE type - we know we need to revamp the encoding attribute to accommodate prime modulus anyway, and we can talk about the scale / delta parameter as well on it
At a higher level (secret dialect) i wonder if we could include a precision loss kind of attribute and auto-calculate the scaling param. Not sure yet though
I'd like to talk at the HEIR WG meeting this week on the LWE type - we know we need to revamp the encoding attribute to accommodate prime modulus anyway, and we can talk about the scale / delta parameter as well on it
At a higher level (secret dialect) i wonder if we could include a precision loss kind of attribute and auto-calculate the scaling param. Not sure yet though
Yep, maybe we can use this as an opportunity to update secret.secret with support for an attribute that describes the plaintext->homomorphic semantics mapping (e.g., going from i32 to $Z_p$, from f32 to some CKKS scaled fixed point representation, etc). Auto-calculating the scaling parameters from some kind of "acceptable precision loss" would be great!
CKKS is well supported now, and we have separate issues for sub-tasks related to scale, etc. raised in this issue. Closing.