Ns/feat/atomic patterns

[nsarlin-zama]

closes: https://github.com/zama-ai/tfhe-rs-internal/issues/643

PR content/description

This PR adds the notion of Atomic Pattern to the shortint layer. This is currently a draft to be able to discuss on the design.

Atomic Pattern

An atomic pattern is a sequence of homomorphic operations that can be executed indefinitely. In TFHE, the standard atomic pattern is the chain of 5 linear operations + KS + PBS. In TFHE-rs, this is implemented at the shortint level by the ServerKey::apply_lookup_table (to be precise this is only the KS/PBS part). The goal of the PR is to add a genericity layer to be able to easily switch atomic pattern without having to rewrite higher level operations.

Implementation

(The names of the trait and types are not definitive)

Currently the shortint ServerKey is represented by this structure:

pub struct ServerKey {
    pub key_switching_key: LweKeyswitchKeyOwned<u64>,
    pub bootstrapping_key: ShortintBootstrappingKey,
    pub message_modulus: MessageModulus,
    pub carry_modulus: CarryModulus,
    pub max_degree: MaxDegree,
    pub max_noise_level: MaxNoiseLevel,
    pub ciphertext_modulus: CiphertextModulus,
    pub pbs_order: PBSOrder,
}

We can split these fields in 2 parts:

• The key materials that are related to the atomic pattern and should use some kind of polymorphism (static or dynamic):
  • key_switching_key, bootstrapping_key and pbs_order
• The metadata that are linked to the shortint encoding

To do that, this PR first adds a trait AtomicPatternOperations. This trait defines the operations that should be supported by all the atomic patterns. It is dyn compatible to allows having atomic patterns as trait objects:

pub trait AtomicPatternOperations {
    fn ciphertext_lwe_dimension(&self) -> LweDimension;

    fn ciphertext_modulus(&self) -> CiphertextModulus;

    fn apply_lookup_table_assign(&self, ct: &mut Ciphertext, acc: &LookupTableOwned);

    fn apply_many_lookup_table(
        &self,
        ct: &Ciphertext,
        lut: &ManyLookupTableOwned,
    ) -> Vec<Ciphertext>;

// ...
}

This trait is first implemented for the "classical" (CJP) atomic pattern:

pub struct ClassicalAtomicPatternServerKey<KeyswitchScalar>
where
    KeyswitchScalar: UnsignedInteger,
{
    pub key_switching_key: LweKeyswitchKeyOwned<KeyswitchScalar>,
    pub bootstrapping_key: ShortintBootstrappingKey,
    pub pbs_order: PBSOrder,
}

From there we have an enum of atomic pattern specific keys that all implement this trait:

pub enum ServerKeyAtomicPattern {
    Classical(ClassicalAtomicPatternServerKey<u64>),
    KeySwitch32(ClassicalAtomicPatternServerKey<u32>),
// and more to come
}

The enum also implements AtomicPatternOperations (the "enum dispatch" design pattern).

Finally, we have the "GenericServerKey" (name not definitive) defined as follow:

pub struct GenericServerKey<AP> {
    pub atomic_pattern: AP,
    pub message_modulus: MessageModulus,
    pub carry_modulus: CarryModulus,
    pub max_degree: MaxDegree,
    pub max_noise_level: MaxNoiseLevel,
    pub ciphertext_modulus: CiphertextModulus,
}

Some type aliases are defined to make it more usable:

pub type ServerKey = GenericServerKey<ServerKeyAtomicPattern>;
pub type ClassicalServerKey = GenericServerKey<ClassicalAtomicPatternServerKey<u64>>;

ServerKey is the one that is used almost everywhere, this reduces the impact on the higher layers. Every methods that use the ServerKey for lookup tables and the shortint encoding are usable without (almost) any modification.

However some features don't fit well in the atomic pattern concept (as I understand it):

• compression
• oprf
• wopbs

For these features, this design allows to create impl blocks that only work for one specific atomic pattern, by using the ClassicalServerKey type. To go from one type to the other, ClassicalServerKey implements TryFrom<ServerKey>.
To make this more efficient, we have 2 "View" types that allow conversions without having to clone the keys:

pub type ServerKeyView<'key> = GenericServerKey<&'key ServerKeyAtomicPattern>;
pub type ClassicalServerKeyView<'key> =
    GenericServerKey<&'key ClassicalAtomicPatternServerKey<u64>>;

In the future, it should be easy to extend the set of supported AP for a feature.
For example we can have an OprfServerKeyAtomicPattern enum with only the subset of ap that support the oprf, and define a type OprfServerKey = GenericServerKey<OprfServerKeyAtomicPattern>;

Check-list:

• Tests for the changes have been added (for bug fixes / features)
• Docs have been added / updated (for bug fixes / features)
• Relevant issues are marked as resolved/closed, related issues are linked in the description
• Check for breaking changes (including serialization changes) and add them to commit message following the conventional commit specification

---

This change is 

[tmontaigu]

Do we foresee that the ServerKeyAtomicPattern is going to have a dynamic variant ServerKeyAtomicPattern::Dynamic(Box<dyn AtomicPattern>) if yes, should the ServerKey just store a Box directly ?

Reviewed 9 of 45 files at r1, all commit messages.
Reviewable status: 9 of 45 files reviewed, 4 unresolved discussions (waiting on @nsarlin-zama)

---

tfhe/src/high_level_api/backward_compatibility/integers.rs line 76 at r1 (raw file):

                                .as_view()
                                .try_into()
                                .map(|sk| old_sk_decompress(sk, a))

We could have let key = sk.key.key.key.as_view().try_into()?; at the beginning of the fn, and use the key in the par_iter.map

Code quote:

                                .try_into()
                                .map(|sk| old_sk_decompress(sk, a))

---

tfhe/src/high_level_api/backward_compatibility/integers.rs line 122 at r1 (raw file):

                                .as_view()
                                .try_into()
                                .map(|sk| old_sk_decompress(sk, a))

Same here, we could do`let key = sk.key.key.key.as_view().try_into()?;

Code quote:

                            sk.key
                                .key
                                .key
                                .as_view()
                                .try_into()
                                .map(|sk| old_sk_decompress(sk, a))

---

tfhe/src/shortint/atomic_pattern.rs line 50 at r1 (raw file):

}

pub trait AtomicPatternMutOperations {

I would have pub trait AtomicPatternMutOperations: AtomicPatternOperations so that having the mut version also gives all the non mul ops

---

tfhe/src/shortint/server_key/mod.rs line 406 at r1 (raw file):

impl<AP: Clone> GenericServerKey<&AP> {
    pub fn into_owned(&self) -> GenericServerKey<AP> {

Generally into_ method take self, so I think a better name would be to_owned

[nsarlin-zama]

Yes the dynamic one is in progress

enum dispatch is supposedly more performant than vtable lookups, and some things are easier to do with enums, like serialization

Reviewable status: 9 of 45 files reviewed, 4 unresolved discussions (waiting on @tmontaigu)

---

tfhe/src/high_level_api/backward_compatibility/integers.rs line 76 at r1 (raw file):

Previously, tmontaigu (tmontaigu) wrote…

We could have let key = sk.key.key.key.as_view().try_into()?; at the beginning of the fn, and use the key in the par_iter.map

good idea !

---

tfhe/src/shortint/atomic_pattern.rs line 50 at r1 (raw file):

Previously, tmontaigu (tmontaigu) wrote…

I would have pub trait AtomicPatternMutOperations: AtomicPatternOperations so that having the mut version also gives all the non mul ops

ok I will do that!

---

tfhe/src/shortint/server_key/mod.rs line 406 at r1 (raw file):

Previously, tmontaigu (tmontaigu) wrote…

Generally into_ method take self, so I think a better name would be to_owned

ok, but to_owned is generally associated with the ToOwned trait, so maybe there is a third option that is less confusing ?