google / heir

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

Adding CKKS to HEIR #808

Open lawrencekhlim opened 4 months ago

lawrencekhlim commented 4 months ago

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

Open Questions

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):

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

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

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.

j2kun commented 4 months ago

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

asraa commented 4 months ago

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

AlexanderViand-Intel commented 4 months ago

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!