Scheme Dialects

[AlexanderViand-Intel]

As discussed in the MLIR Dialects Session, I'm setting up three discussions topics, corresponding to the three streams identified.
This is the Scheme Dialects discussion, you can find the High-Level FHE Dialect and Poly/Math Dialects discussions at the links.

The goal of this abstraction/dialect is to provide a dialect (or set of dialects) that can act as an intermediate representation and exchange format for FHE computations at the "circuit"-ish level.

Technical/MLIR Challenges

As this is supposed to be something of an "XONN/StableHLO for FHE", we should strive to follow modern MLIR best practices in dialect design. However, as the MLIR project has been continuously evolving, it is not immediately clear what exactly these best practices currently are and which, if any, of the upstream dialects showcase them.
In addition to providing the "raw" dialect(s), the aim of this project is to "include batteries", i.e. features and tooling around the dialect(s). This includes pretty printers for concise textual representations, efficient binarization to allow this to be used as an actual interchange format, and a yet-to-be-determined set of canonicalizations, folds, and potentially passes.

Conceptual Challenges

FHE schemes tend to define a natural "API" of sorts, which is reflected in the existing FHE libraries. Therefore, it should be relatively straightforward to define per-scheme dialects, with the majority of the challenge lying in ensuring that the dialects are expressive enough to support a variety of different implementations (potentially as drastically different as RNS vs non-RNS implementations, but certainly differing in approaches to noise management during ciphertext maintenance operations). From this starting point, we should then try to identify common components that can be factored out, enabling a slow transition to a universal dialect.

Next steps

The presence of high-quality dialects to represent FHE computations at this level appears to be a blocking issue for (parts of) the further development of the MLIR version of the FHE Transpiler and new versions of HECO. As a result, the goal is to start rapidly prototyping working dialects for the schemes/variants that are most relevant for these uses cases. The timeline for this is 1 month, at which point we should report our progress back to the working group and switch to exploring towards a universal abstraction.

Tasks

• Please register your interest in actively contributing to this abstraction/dialect by commenting below, and suggest possible next task.
• If interested, please check whether you are able to agree to the Google CLA, as the current idea is to have these dialects live in the HEIR directory, with upstreaming not an immediate goal (upstreaming should become more viable if a "universal" dialect emerges)

[j2kun]

I am happy to prototype a CGGI dialect, since that's the scheme I'm most familiar with. I would also like to review the CKKS and BGV-like dialect drafts so I can better learn them myself.

[j2kun]

Started a prototype in #119

[AlexanderViand-Intel]

In yesterday's Working Group meeting (see notes), the desire to have a target/hardware-agnostic interface (e.g., as a specification for accelerators to implement) came up.

While scheme-specific dialects (like Zama's tfhe-rs-derived CGGI one, or HECO's seal-derived BFV one) seems like an ideal candidate for this at first glance, programs at this abstraction level are usually already (over)specified with a specific implementation in mind. These dialects might already encode parameter and algorithm choices (even if only partially, i.e., specifying number of RNS limbs, coeff vs eval form, etc) that should be chosen to fit the specific target implementation. The following is a brief list of examples of choices that might need to be made in a target-dependent way.

• Increase bootstrapping precision by increasing parameters, or by using a hierachy of multiple lower-precision bootstraps?
• Compute directly at large bit-widths (e.g., 128) or compose from smaller units (e.g, 32)
• (CGGI-specific?) Whether to use FFT-friendly parameters or whether to use only NTT
• (CGGI-specific?) Using BGV-style multiplication, or LUT-based multiplication
• ... more examples certainly exist, so please reply in the comments if you can think of some!

Therefore, I propose to have an "intermediate" abstraction step between high-level and (traditional) scheme-specific, something we might call semantic/declarative scheme-level IR. This would omit details such as parameters/etc, and instead express requirements (e.g., precision of arithmetic operations, precision of bootstrap, etc, all in bits) similar to the high-level dialect. In order to be a meaningfully "lower" level abstraction, the intent is still to restrict the set of operations to those that have "direct" equivalents in the underlying scheme. This could then serve as a natural separation point between more "high-level" oriented compilers (e.g. open-source HECO) and more "low-level" oriented compilers (e.g., toolchains for DPRIVE accelerators). In addition, it would be a very useful interchange format for "apples to apples" benchmarking between different libraries/accelerators.

Some current open questions are:

• Do we need to specify vector/ctxt size for RLWE schemes? While there are techniques to bridge gaps between vector and ctxt sizes (c.f. HElayers, EVA Improved) they feel more high-level than this, require significant engineering effort to implement, and might have negative impacts on performance.
• Should these operations live in their own MLIR dialect? If not, should they be separate operations or simply "overloads" of the normal ones?

[j2kun]

I feel like even the LUT size is an optimizer choice. By LUT-size I mean the bit width of the semantic input/outputs. E.g., a 2-LUT would be equivalent to binary boolean gates, and a 4-LUT would have 4-bit input/output.

From my experience with tfhe-rs and my own LUT-based implementations, two backends can have very different performance even for the same LUT size. For example, tfhe-rs has a dramatic parameter increase from 3-LUT to 4-LUT, making 3-LUT more performant in most of my experiments (that, admittedly, focus on circuit optimization, not high-level decomposition of large-bit-width LUTs into a hierarchy of smaller bit-width LUTs). But our implementation (now quietly public in https://github.com/google/jaxite before an upcoming announcement) has a smaller performance difference between 3-LUT and 4-LUT, so 4-LUT can a better choice if it can make the circuit much smaller or more parallelizable. Part of the reason for this discrepancy, I believe, is that Zama uses 64-bit integers for the underlying types while we use 32-bit integers, and Zama has a different choice of RLWE and decomposition parameters (perhaps due to lower failure tolerances? or due to less frequent bootstrapping?).

[j2kun]

And definitely to add to the "FFT vs NTT friendly parameters", on TPU I think we would have neither, so the parameters would likely be chosen for how well the resulting matrices fit into the matrix multiplication unit on the particular TPU generation in question (all of which have different high-bandwidth memory sizes)

[FlorentCLMichel]

Your point about the LUT size is a very good one! And thanks for the link to Jaxite! It looks super interesting; looking forward to the public announcement!

[AlexanderViand-Intel]

Some very good points being raised! And congratulations on getting your backend open-sourced 👍

I figured it might be helpful to have a few end-to-end examples to show how these things arise and what the possibilities are. Here's my attempt at one for now:

High-level code for a comparison in some frontend:

# assume a, b are secret integers with 16 bits but we need only 8 bit precision here
c = a < b  

Gets translated to high-level dialect:

%c = fhe.cmp(%a, %b) {prec = 8 } : (!fhe.secret<i16>,!fhe.secret<i16>) -> !fhe.secret<bool>

At this point, this could be translated to (a) bit-wise encryption in CGGI or an RLWE scheme, (b) polynomial interpolation in an RLWE scheme, or (c) a LUT based comparison in CGGI, and we'll do the latter here.

Note, that this lowering step could still use the fhe dialect in an intermediate step (going from a<b to sign(a-b) ), which might be a common step between different approaches (e.g., lots of interpolation based things do this, too). Not sure if that's a good idea, though:

%t = fhe.sub(%a, %b) {prec = 8 } : (!fhe.secret<i16>,!fhe.secret<i16>) -> !fhe.secret<i16>
%c = fhe.sign(%t) {prec = 8 } : (!fhe.secret<i16>) -> !fhe.secret<bool>

Whether or not we have the intermediate step, we'd arrive at something along these lines, where the * indicates this intermediate abstraction level (Here, I'm the least sure about if this is indeed what it should look like.):

%t = cggi.sub*(%a, %b) : (!cggi.ctxt*<8>,!cggi.ctxt*<8>) -> !cggi.ctxt*<8>
%c = cggi.lut*(%t) {[-128 : 1, -127 : 1, .... , 0 : 0, 1 : 0, .... , 127:0]} :  (!cggi.ctxt*<8>,!cggi.ctxt*<8>) -> !cggi.ctxt*<1>

This would then be either lowered to something with (very) large parameters (not sure where exactly that info would go and what it would look like) and a single LUT:

%t = cggi.sub(%a, %b) : (!cggi.ctxt<LARGE_PARAM_INFO>,!cggi.ctxt<LARGE_PARAM_INFO>) -> !cggi.ctxt<LARGE_PARAM_INFO>
%c = cggi.lut(%t) {[-128 : 1, -127 : 1, .... , 0 : 0, 1 : 0, .... , 127:0]} :  (!cggi.ctxt<LARGE_PARAM_INFO>,!cggi.ctxt<LARGE_PARAM_INFO>) -> !cggi.ctxt<LARGE_PARAM_INFO>

or into something with smaller parameters and multiple LUTs, but I don't think I'm familiar enough with the multi-LUT approach to write something that could pass for valid IR for that ;)

I hope this little exercise helps a bit in visualizing the different abstraction levels. I think it'd be very interesting to repeat this exercise (not just specifically for this abstraction, but in general) with more than just a computation snippet, i.e. also include the function, input/outputs/keys/etc, to really understand if there's something we've missed and to find the answers to decisions such as "Where should the security level be specified and how"? (I omitted this in my example above, because I think doing this on a per-operation basis might be odd? I also omitted key identifiers in the high-level dialect, but maybe keys should have required security levels associated with them?)

[FlorentCLMichel]

Thanks for that! I think having a few end-to-end examples will be very useful indeed!
I am not familiar enough with the LLVM IR syntax to comment on the snippets, but the logic and what I can understand from the snippets seem very sensible to me.

Regarding the security level, would it make sense to have it as a parameter (not sure if that's the right word) to !fhe.secret along with i16, and could this potentially be a way to have a single number in the high-level dialect yielding different parameters (like the RLWE size and modulus) when lowering (potentially in a hardware-dependent way)?
(Not sure whether it would be a good solution, though. I agree that it would be good to have more examples and/or a larger one to get a broader view of the process. I'll try to think of something.)

[FlorentCLMichel]

Here is an attempt at (sketch of an) example for integer division. (I'll mostly write pseudo-code as I'm not quite confident enough in my MLIR knowledge.)

High-level code:

# a and b are secret 16-bits unsigned integers; with want a correct result up to 0.5%
c = a / b;

Translation into high-level dialect:

%c = fhe.uint_div(%a, %b) {prec = 8} : (!fhe.secret<u16>,!fhe.secret<u16>) -> !fhe.secret<u16>

(Open question 1: Do we want to return a special value indicating an error in case b = 0? In the following, I assume for simplicity that the algorithm used for division deals with that case in some ‘appropriate’ way.)

There are multiple ways this could be lowered, e.g.:

• Evaluation of a a huge LUT in CGGI (probably impractical in this case, but may become feasible for 8-bit numbers).
• Evaluation of a sequence of LUTs in CGGI.
• Evaluation of a boolean circuit in CGGI.
• Evaluation using Goldschmidt's algorithm in CGGI, BGV, B/FV, or CKKS.
• Polynomial approximation using BGV, B/FV, or CKKS.

(Open question 2: At which point is the selection of the best algorithm and scheme made?)

For definiteness, let's assume we choose to use Goldschmidt's algorithm implemented with CKKS. We then need to

• Define the number of iterations N (may be hard-coded for a given target accuracy).
• Define a power-of-2 scale factor s.
• Define a starting guess for s / b (which does not actually need to depend on b; in practice, if b is guaranteed to be in a given range which does not include 0, one can pick a fixed value which works over the whole range).
• Define starting values for u, v, and w as: u = a, v = b, and w as the approximation of s / b.
• At each iteration, perform the following operations:
  • u <- (u * w) / s
  • v <- (v * w) / s
  • w <- s + 1 - v
• Return u.

The multiplications, division by s (presumably implemented as homomorphic bit-shifting), and subtraction should then be lowered to, e.g, ckks.mul, ckks.shr, and ckks.sub.

I realise this is a bit sketchy; but hopefully it helps showing a possible sequence of operations on another example.

[asraa]

(Open question 2: At which point is the selection of the best algorithm and scheme made?)

For a decision on a particular algorithm, if we lower to a general "CGGI" dialect like options 1-4, optimization passes can convert between the options depending on the whole IR and other parameters. Would it need to happen at this step even though it may depend on lower-level design choices? e.g. what if the target CGGI library's large LUT is somehow very fast and the small LUT is very slow - we'd want to select option 1 instead of 2: how can we make that choice since it's based on a lower level option? Or do we wait to lower down to the target library dialect and figure out a way to combine LUTs (that seems harder?)

[j2kun]

I'm seeing the prototypes that @asraa put together for BGV and the one I'm working on for CGGI (#119), and I was thinking we might want to have a common type-only dialect for LWE and RLWE (or combine them into one "generalized LWE" type a la Zama) so that ciphertext types can be shared across scheme-specific dialects.

The BGV type has a ring array attr, and the RLWE type for CGGI would just need a single ring attribute, so maybe that could just be an array of size 1 and the CGGI op verifiers can check that. I think only CGGI uses vanilla LWE, but it might still be nice to have it as a separate concept.

Thoughts?

[mr0re1]

@j2kun what's your thoughts on doing a single GLWE , instead of LWE and RLWE separately?

[AlexanderViand-Intel]

@j2kun what's your thoughts on doing a single GLWE , instead of LWE and RLWE separately?

Expanding on this a bit: In the last meeting, we discussed the need for a scheme-independent "toolbox" dialect that houses common algorithmic building blocks of (R)LWE based schemes that are neither fully scheme specific nor are generic enough for the underlying dialects (i.e., linalg or poly). For example, RLWE-related things in this toolbox dialect might include digit decomposition, fast basis conversion, and other "primitives" that show up as part of, e.g., keyswitching. Keyswitching itself would, however, probably end up in the scheme-specific dialects as it requires concepts such as kleys, ctxt, etc, whereas the idea for the toolbox dialect would be to operate on (tensors of) polynomials and LWE samples (maybe also just represented as tensors?) EDIT: It's "cheap" to move operations between dialects (mostly search&replace based refactoring) so we can always postpone decisions such as where keyswitch goes until we've got more of the lowerings and can see where they fit.

[j2kun]

Moving to #147

[j2kun]

I'm starting to second guess the idea that the CGGI dialect should have its types be aliases to tensor<sample_size, encoding>. In particular, that only really makes sense if CGGI is going to lower through the rest of MLIR (and poly) to LLVM or some hardware dialect. If, as we intend, we want to be able to export CGGI to LLVM + API calls against an existing FHE library, then it makes more sense to have a somewhat agnostic CGGI ciphertext type that can lower to the library dialect's ciphertext type (which may be opaque to the library user), or to a tensor<whatever>, but have the lowering decide that instead of forcing it on all users of the dialect.

[BourgerieQuentin]

Yes make sense to be an abstract type and the "materialization" to tensor or to something else happens when we lower, else it make a dependency between the abstraction and the actual implementation 👍