Secret dialect

[j2kun]

Putting aside a separate thread (from #53 (reply in thread)) to discuss aspects of a high-level secret dialect (even higher level than something like "encrypted arith").

A draft PR showing the basic structure is #78

Some example MLIR might look like

%X = arith.constant 7 : i32
%Y = secret.conceal %value : i32 -> !secret.secret<i32>
%Z = secret.generic {library_name = "foo"}
  ins(%X, %Y : i32, !secret.secret<i32>) {
  ^bb0(%x: i32, %y: i32) :
    %d = arith.addi %x, %y: i32
    secret.yield %d : i32
  } -> (!secret.secret<i32>)
%W = secret.reveal %Z : !secret.secret<i32> -> i32

The motivation behind this dialect is:

• As a front-end dialect, compatible with any MLIR dialect.
• To preserve high-level op semantics (in the style of linalg.generic) so that we can lower targeted ops to specially handled ops in the HEIR ecosystem.
• To enable some early transformations that are unrelated to FHE, such as folding out plaintext computations or passing through loops with static bounds.

The central op here is secret.generic which lifts a plaintext computation to operate on one or more secrets, with secret.secret being a wrapper type around an arbitrary MLIR type.

In the last MLIR meeting @AlexanderViand-Intel expressed that it would be too difficult to make useful optimization decisions at this level, requiring original research to develop a good cost model at least. I think it would still have value as an easy front-end to target, and doesn't prevent users from targeting lower-level dialects like encrypted_arith/math/linalg if they prefer.

[BourgerieQuentin]

This looks good to me as it look like we can express anything with it and at the same time keep high level information (here the library_name, btw I don't like "library_name" maybe "op_name" or something like that may fit more what we want to express but it's a detail).

I'm ok with @AlexanderViand it look like difficult at this level to make useful optimization, but I guess is not the purpose of this dialect as it should be designed to be the most generic one in order to be lowered to scheme specific dialect. Maybe we could have a collection of transformation at this level to offer basic transformation, like eliminate addition with 0, multiplication by 0/1, merge consecutive table lookup etc... About the implementation of such simplification I'm not sure it should be folder/canonicalizer or separate rewrite patterns, as maybe some simplification could be suboptimal for a specific scheme, thus this would be the responsability of compilers to select patterns that make sense for it specific scheme?

I think we should also have a collection of non-generic operators that make codegen from frontend more easy to do, and a pass that convert non-generic operator into secret.generic operator (just like linalg named ops). One of the question that raise is where is the boundary of the secret named operator? To take as example on Concrete compiler we have very high level fhe operators that works on tensors with a broadcasting semantics that is usefull for our frontend codegen, but is that make sense to have this operator in common?

Should we have a way to express that a secret is under a certain key (maybe something optional ?), if yes is that means we should also have a kind of keyswitch operator at this level? And what about public encryption, threshold decryption?

[j2kun]

So all the usual MLIR passes should work inside secret.generic at this level without us doing anything, so long as the ops being optimized are inside a single region. So I think the initial passes would be something like this:

1. Using whatever front-end specific mechanism available, designate which entry values are private (say, an annotation on a function argument or declared variable).
2. Write a secret pass that threads secret.secret through a program, lifting designated values to secret.secret<T>, and lift all blocks that touch a secret into a secret.generic op (e.g., one block is one big secret.generic with lots of cleartext ops inside).
3. Run normal MLIR passes on the output of (2), which would give all the folding and canonicalizations we want.
4. Run a secret pass that splits up secret.generic into generics where the blocks are individual ops (folding any generic's that have no secrets into plain MLIR ops
5. Lower to something like fhe_math.generic or fhe_arith, possibly preserving the name of the original op so that later passes/lowerings can be aware of the semantics of the op. This would be useful for a case like, say, a tosa.relu operation that is lowered to fhe_math.generic, the body of which has an implementation of relu using supported ops, but it retains a library_name attribute so that hardware with a special fhe_relu support can lower the entire generic op in a special way.

In this sense, I think secret is the easiest possible entry point for a front-end language, and then the compiler is responsible for managing the lowering through more fhe-aware dialects with named or generic ops.

But to your last point, secret definitely should have the concept of which key the secret is encrypted for, I just wanted to keep the initial PR small so that we could get something established and iterate on it. I just filed #122 for adding an attribute to specify which key. I'm not sure about threshold encryption or public key encryption (probably public key encryption can be the same implementation as symmetric key at secret, and just use an opaque identifier to represent a key-pair).