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 AtomicPattern. 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 AtomicPattern{
    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
{
    pub key_switching_key: LweKeyswitchKeyOwned<u64>,
    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),
    KeySwitch32(KS32AtomicPatternServerKey),
// and more to come
}

The enum also implements AtomicPattern (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 in the upper layers, this reduces the breaking changes. Every methods that use the ServerKey for lookup tables and the shortint encoding are usable without (almost) any modification.

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

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>;

Dynamic AP

The GenericServerKey enum has a last specific variant:

pub enum AtomicPatternServerKey {
    Classical(ClassicalAtomicPatternServerKey),
    KeySwitch32(KS32AtomicPatternServerKey),
    #[serde(skip)]
    Dynamic(Box<dyn private::DynamicAtomicPattern>),
}

This allows to create a server key for an AP defined by the user. the trait DynamicAtomicPattern is sealed but automatically impl for most of the types that implement AtomicPattern. This trait allows to impl Clone and PartialEq for the dynamic case.

KS32

This PR also implements the DP-KS-PBS-AP-with-32-bit-KS-modulus atomic pattern to have an example of what can be done using this structures.

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 ?

[IceTDrinker]

After discussing I agree, no need to have a special handling for the new generic parameter as it's "just" a result of "contamination" from an inner member that handles its own upgrade automatically

Just one thing around a parameter alias that I would prefer not done this way, otherwise, this is the last round for this PR

Reviewed 1 of 1 files at r9, 11 of 14 files at r10, 9 of 10 files at r11, 2 of 2 files at r12, 1 of 1 files at r13, all commit messages.
Reviewable status: all files reviewed, 13 unresolved discussions (waiting on @mayeul-zama and @nsarlin-zama)

---

tfhe/src/shortint/atomic_pattern/ks32.rs line 66 at r6 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

The problem here is that the client key is not generic, it stores its keys using u64 Scalar.
I think at some point we would have to make the ClientKey an enum with AP specific variants like the ServerKey, this is a workaround to avoid that.

understood but having additional copies of the secret key in memory is something we don't want, please have a follow up issues, this can either be managed as you say or making the core APIs even more flexible with a scalar only tied to the ciphertext modulus and compatible scalars for both input and output keys

---

tfhe/src/shortint/atomic_pattern/mod.rs line 547 at r6 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

done ! I forgot to update this match case

do we want a wildcard for all other cases or rather a properly exhaustive code so that this is not forgotten later ? follow up issue or in the compressed key follow up

---

tfhe/src/shortint/server_key/modulus_switched_compression.rs line 33 at r6 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

If I revert this this lint triggers: https://rust-lang.github.io/rust-clippy/master/index.html#option_if_let_else. For me it's more a question of habit so I don't have any strong opinion.

Just that I would keep the compression which is common to both after the application of MS noise reduction or not

---

tfhe/src/shortint/atomic_pattern/classical.rs line 33 at r6 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

this was pub when it was in the server key so I simply reused the previous pattern. Do you think I should make them private ?

long term yes, but it will do for now

---

tfhe/src/shortint/atomic_pattern/mod.rs line 557 at r13 (raw file):

#[cfg(test)]
mod test {
    use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128;

create a test_params alias instead please, for now we are not adverstising KS32 until it's prod ready

in shortint::parameters::test_params I believe

---

tfhe/src/shortint/parameters/aliases.rs line 174 at r13 (raw file):

// KS32 AP
pub const PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128: KeySwitch32PBSParameters =

as said elsewhere I would prefer a test_params alias

---

tfhe/src/core_crypto/entities/compressed_modulus_switched_multi_bit_lwe_ciphertext.rs line 193 at r12 (raw file):

    ) -> Self {
        assert_eq!(
            packed_mask.initial_len() % grouping_factor.0,

careful if grouping factor == 0 (which should not happen, but you know), let's add it as a follow up in the compressed server key PR as a separate commit to avoid delaying this one more

[nsarlin-zama]

Reviewable status: all files reviewed, 10 unresolved discussions (waiting on @IceTDrinker and @mayeul-zama)

---

tfhe/src/core_crypto/entities/compressed_modulus_switched_multi_bit_lwe_ciphertext.rs line 193 at r12 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

careful if grouping factor == 0 (which should not happen, but you know), let's add it as a follow up in the compressed server key PR as a separate commit to avoid delaying this one more

done
Added an assert_ne!(grouping_factor, 0)

---

tfhe/src/shortint/atomic_pattern/ks32.rs line 66 at r6 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

understood but having additional copies of the secret key in memory is something we don't want, please have a follow up issues, this can either be managed as you say or making the core APIs even more flexible with a scalar only tied to the ciphertext modulus and compatible scalars for both input and output keys

https://github.com/zama-ai/tfhe-rs-internal/issues/1007#issue-3028543196

---

tfhe/src/shortint/atomic_pattern/mod.rs line 547 at r6 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

do we want a wildcard for all other cases or rather a properly exhaustive code so that this is not forgotten later ? follow up issue or in the compressed key follow up

the problem is that having all the other cases exhaustively handled does not scale well as we need to catch all possible combinations.

---

tfhe/src/shortint/atomic_pattern/mod.rs line 557 at r13 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

create a test_params alias instead please, for now we are not adverstising KS32 until it's prod ready

in shortint::parameters::test_params I believe

done

---

tfhe/src/shortint/parameters/aliases.rs line 174 at r13 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

as said elsewhere I would prefer a test_params alias

done

---

tfhe/src/shortint/server_key/modulus_switched_compression.rs line 33 at r6 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

Just that I would keep the compression which is common to both after the application of MS noise reduction or not

The problem here is that in one branch I have LweCiphertext<&[Scalar]> and in the other one LweCiphertext<Vec<Scalar>> so I cannot return them both in the branches.

I cannot declare the input_improved_before_ms uninitialized before the block because of the closure. I need to intialize it with an empty ciphertext which feels weird.

I think I here the 2 options are to disable the lint or keep it like that

[IceTDrinker]

Huge effort on your part @nsarlin-zama a big big thank you for this work, I know this has not been easy (and a long time coming) and the reviews were thorough

Let's go!

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

---

tfhe/src/shortint/atomic_pattern/mod.rs line 547 at r6 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

the problem is that having all the other cases exhaustively handled does not scale well as we need to catch all possible combinations.

ah right, misread this :/

---

tfhe/src/shortint/server_key/modulus_switched_compression.rs line 33 at r6 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

The problem here is that in one branch I have LweCiphertext<&[Scalar]> and in the other one LweCiphertext<Vec<Scalar>> so I cannot return them both in the branches.

I cannot declare the input_improved_before_ms uninitialized before the block because of the closure. I need to intialize it with an empty ciphertext which feels weird.

I think I here the 2 options are to disable the lint or keep it like that

Potential solutions (but let's go ahead with the code as it is now): .as_view to have the same return type (that's what was done before the change), as for the unitialized variable I know you can do some "forward declarations" with "just" a let x; and then as long as the compiler can prove that x is initialized in the branch where it is used it does not bother you, but I really don't know how that behaves with closures

[nsarlin-zama]

Reviewable status: all files reviewed, 9 unresolved discussions (waiting on @IceTDrinker and @mayeul-zama)

---

tfhe/src/shortint/server_key/modulus_switched_compression.rs line 33 at r6 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

Potential solutions (but let's go ahead with the code as it is now): .as_view to have the same return type (that's what was done before the change), as for the unitialized variable I know you can do some "forward declarations" with "just" a let x; and then as long as the compiler can prove that x is initialized in the branch where it is used it does not bother you, but I really don't know how that behaves with closures

No that's what I was trying but it does not work with the closure, here it needs to be initialized