feat(shortint): adds generic client key for atomic pattern support

[nsarlin-zama]

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

PR content/description

This removes the unneeded lwe key copies for the KS32. The chosen solution is to make the client key dependent on the atomic pattern.

The implementation follows the same structure as the implementation of the server key atomic pattern.

Here, the different client keys implement a new trait, EncryptionAtomicPattern:

pub trait EncryptionAtomicPattern {
    /// The parameters associated with this client key
    fn parameters(&self) -> ShortintParameterSet;

    /// The secret key used for encryption
    fn encryption_key(&self) -> LweSecretKeyView<'_, u64>;

    /// The noise distribution used for encryption
    fn encryption_noise(&self) -> DynamicDistribution<u64>;

    /// The kind of atomic pattern that will be used by the generated
    /// [`ServerKey`](crate::shortint::ServerKey)
    fn kind(&self) -> AtomicPatternKind;

    /// The secret key used to generate a KeySwitchingKey for these parameters
    fn keyswitch_encryption_key(
        &self,
        params: ShortintKeySwitchingParameters,
    ) -> LweSecretKeyView<'_, u64>;

    /// The noise distribution used to generate a KeySwitchingKey for these parameters
    fn keyswitch_encryption_noise(
        &self,
        params: ShortintKeySwitchingParameters,
    ) -> DynamicDistribution<u64>;

This trait is roughly used to give access to the encryption keys, so that the amount of code we have to write for a new AP is minimal, and we can reuse all the different encryption functions.

Similarly to what has been done with the ServerKey, the ClientKey is now converted to a GenericClientKey<AP> with an enum ClientKeyAtomicPattern that implement this trait for the default supported AP (standard and KS32).

At the moment, no support for the dynamic AP has been implemented on the client key.

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 

[IceTDrinker]

Some comments already, the goal is achieved, there is something unclear at the moment with the keyswitching stuff, we would not necessarily limit ourselves to keyswitching to the big key, given that keyswitching to the small key means smaller keys and less noise potentially (and cheaper computations)

Reviewed 29 of 63 files at r1, 3 of 5 files at r2, all commit messages.
Reviewable status: 31 of 64 files reviewed, 5 unresolved discussions (waiting on @mayeul-zama and @tmontaigu)

---

tfhe/src/shortint/client_key/atomic_pattern/ks32.rs line 168 at r2 (raw file):

    }

    fn keyswitch_encryption_key(

not sure this is right, the ksk is encrypted under the small key

---

tfhe/src/shortint/client_key/atomic_pattern/ks32.rs line 174 at r2 (raw file):

        match params.destination_key {
            EncryptionKeyChoice::Big => self.large_lwe_secret_key(),
            EncryptionKeyChoice::Small => {

we could imagine a world where we ks to the small key unfortunately

---

tfhe/src/shortint/client_key/atomic_pattern/ks32.rs line 180 at r2 (raw file):

    }

    fn keyswitch_encryption_noise(

same here something does not seem to match

---

tfhe/src/shortint/backward_compatibility/client_key/mod.rs line 46 at r1 (raw file):

                .downcast_ref::<GenericClientKey<AP>>()
                .unwrap() // We know from the TypeId that AP is of the right type so we can unwrap
                .clone())

leaves some secret keys copies, do we have to clone ?

---

tfhe/src/shortint/backward_compatibility/client_key/mod.rs line 54 at r1 (raw file):

                .downcast_ref::<GenericClientKey<AP>>()
                .unwrap() // We know from the TypeId that AP is of the right type so we can unwrap
                .clone())

leaves some secret keys copies, do we have to clone ?

[tmontaigu]

tfhe/src/shortint/backward_compatibility/client_key/mod.rs line 46 at r1 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

leaves some secret keys copies, do we have to clone ?

Can't we have two separate impls ?

impl Upgrade<GenericClientKey<AtomicPatternClientKey>> for ClientKeyV0 {
    type Error = Error;

    fn upgrade(self) -> Result<GenericClientKey<AtomicPatternClientKey>, Self::Error> {
        let ap_params = self.parameters.pbs_parameters().ok_or_else(|| {
            Error::new(
                "ClientKey from TFHE-rs 1.2 and before needs PBS parameters to be upgraded to the latest version"
                    .to_string(),
            )
        })?;

        let std_ap = StandardAtomicPatternClientKey::from_raw_parts(
            self.glwe_secret_key,
            self.lwe_secret_key,
            ap_params,
            self.parameters.wopbs_parameters(),
        );

        let atomic_pattern = AtomicPatternClientKey::Standard(std_ap);
        Ok(ClientKey { atomic_pattern })
    }
}

impl Upgrade<GenericClientKey<AtomicPatternClientKey>> for ClientKeyV0 {...}

[nsarlin-zama]

Reviewable status: 31 of 64 files reviewed, 3 unresolved discussions (waiting on @IceTDrinker, @mayeul-zama, and @tmontaigu)

---

tfhe/src/shortint/backward_compatibility/client_key/mod.rs line 46 at r1 (raw file):

Previously, tmontaigu (tmontaigu) wrote…

Can't we have two separate impls ?

impl Upgrade<GenericClientKey<AtomicPatternClientKey>> for ClientKeyV0 {
    type Error = Error;

    fn upgrade(self) -> Result<GenericClientKey<AtomicPatternClientKey>, Self::Error> {
        let ap_params = self.parameters.pbs_parameters().ok_or_else(|| {
            Error::new(
                "ClientKey from TFHE-rs 1.2 and before needs PBS parameters to be upgraded to the latest version"
                    .to_string(),
            )
        })?;

        let std_ap = StandardAtomicPatternClientKey::from_raw_parts(
            self.glwe_secret_key,
            self.lwe_secret_key,
            ap_params,
            self.parameters.wopbs_parameters(),
        );

        let atomic_pattern = AtomicPatternClientKey::Standard(std_ap);
        Ok(ClientKey { atomic_pattern })
    }
}

impl Upgrade<GenericClientKey<AtomicPatternClientKey>> for ClientKeyV0 {...}

Good idea ! I also fixed the server key which used the same pattern

---

tfhe/src/shortint/backward_compatibility/client_key/mod.rs line 54 at r1 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

leaves some secret keys copies, do we have to clone ?

done

---

tfhe/src/shortint/client_key/atomic_pattern/ks32.rs line 168 at r2 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

not sure this is right, the ksk is encrypted under the small key

I removed this impl, the ksk now only supports the std atomic pattern

---

tfhe/src/shortint/client_key/atomic_pattern/ks32.rs line 174 at r2 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

we could imagine a world where we ks to the small key unfortunately

removed

---

tfhe/src/shortint/client_key/atomic_pattern/ks32.rs line 180 at r2 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

same here something does not seem to match

removed

[IceTDrinker]

Small nits, otherwise looks good !

Reviewed 30 of 63 files at r1, 9 of 9 files at r3, all commit messages.
Reviewable status: all files reviewed, 2 unresolved discussions (waiting on @mayeul-zama and @tmontaigu)

---

tfhe/src/shortint/client_key/atomic_pattern/mod.rs line 41 at r3 (raw file):

}

// This blancket impl is used to allow "views" of server keys, without having to re-implement the

of client keys

---

tfhe/src/shortint/backward_compatibility/server_key/mod.rs line 105 at r3 (raw file):

}

impl Upgrade<StandardServerKey> for ServerKeyV1 {

is the modulus switch noise reduction key also concerned by this ? I feel like it may not be the case as we need an impl for all possible Scalars for the generics but I could be wrong.

Does not need to be done in this PR as it's unrelated

[nsarlin-zama]

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

---

tfhe/src/shortint/backward_compatibility/server_key/mod.rs line 105 at r3 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

is the modulus switch noise reduction key also concerned by this ? I feel like it may not be the case as we need an impl for all possible Scalars for the generics but I could be wrong.

Does not need to be done in this PR as it's unrelated

Ok, so it reminds me why I used the any, this kind of also applies to the client/server keys so we might need to reconsider (but it is less impacting here thanks to the AtomicPatternServerKey enum).

If we have a type with V0 without generics and V1 with a generic:

#[derive(Version)]
struct MyTypeV0(u64);

#[derive(Version)]
struct MyTypeV1<T>(T);

impl Upgrade<MyTypeV1<u64>> for MyTypeV0 {
  /// ...
}

#[derive(VersionsDispatch)]
enum MyTypeVersions<T> {
  V0(MyTypeV0),
  V1(MyTypeV1<T>)
}

If we define only the upgrades between V0 to V1 for specific statically defined T it might block us in the future. For example if later we have a V2, it means that we will only be able to go from V1 to V2 for this specific T.

#[derive(Version)]
struct MyTypeV0(u64);

#[derive(Version)]
struct MyTypeV1<T>(T);

impl Upgrade<MyTypeV1<u64>> for MyTypeV0 {
  /// ...
}

#[derive(Version)]
struct MyTypeV2<T> {

}

impl<T> Upgrade<MyTypeV2<T>> for MyTypeV1<T> {
  /// ...
}

#[derive(VersionsDispatch)]
enum MyTypeVersions<T> {
  V0(MyTypeV0),
  V1(MyTypeV1<T>),
  V2(MyTypeV2<T>)
}

Here even if the upgrade from V1 to V2 is generic, the unversionize trait will only be derived if T is u64 because it needs to be able to run the full upgrade chain. I had to do dynamic type check because that way we can support going from V0 to V1 for u64 and from V1 to V2 for any T.

[nsarlin-zama]

I reverted to the dynamic type check during the upgrade but this time I used Box::downcast instead of downcast_ref, which removes the need to clone

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

[IceTDrinker]

Last questions

Reviewed 5 of 5 files at r4, all commit messages.
Reviewable status: all files reviewed, 3 unresolved discussions (waiting on @IceTDrinker, @mayeul-zama, @nsarlin-zama, and @tmontaigu)

---

tfhe/src/shortint/backward_compatibility/server_key/mod.rs line 105 at r3 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

Ok, so it reminds me why I used the any, this kind of also applies to the client/server keys so we might need to reconsider (but it is less impacting here thanks to the AtomicPatternServerKey enum).

If we have a type with V0 without generics and V1 with a generic (example without the macros to reduce the boilerplate):

#[derive(Version)]
struct MyTypeV0(u64);

#[derive(Version)]
struct MyTypeV1<T>(T);

impl Upgrade<MyTypeV1<u64>> for MyTypeV0 {
  /// ...
}

#[derive(VersionsDispatch)]
enum MyTypeVersions<T> {
  V0(MyTypeV0),
  V1(MyTypeV1<T>)
}

If we define only the upgrades between V0 to V1 for specific statically defined T it might block us in the future. For example if later we have a V2, it means that we will only be able to go from V1 to V2 for this specific T.

#[derive(Version)]
struct MyTypeV0(u64);

#[derive(Version)]
struct MyTypeV1<T>(T);

impl Upgrade<MyTypeV1<u64>> for MyTypeV0 {
  /// ...
}

#[derive(Version)]
struct MyTypeV2<T> {

}

impl<T> Upgrade<MyTypeV2<T>> for MyTypeV1<T> {
  /// ...
}

#[derive(VersionsDispatch)]
enum MyTypeVersions<T> {
  V0(MyTypeV0),
  V1(MyTypeV1<T>),
  V2(MyTypeV2<T>)
}

Here even if the upgrade from V1 to V2 is generic, the unversionize trait will only be derived if T is u64 because it needs to be able to run the full upgrade chain. I had to do dynamic type check because that way we can support going from V0 to V1 for u64 and from V1 to V2 for any T.

Indeed

---

tfhe/src/shortint/backward_compatibility/client_key/mod.rs line 55 at r4 (raw file):

        } else {
            Err(Error::new(
                "ClientKey from TFHE-rs 1.2 and before can only be deserialized to the classical \

standard atomic pattern ?

---

tfhe/src/shortint/backward_compatibility/server_key/mod.rs line 115 at r4 (raw file):

                self.max_noise_level,
            ));
            Ok(*sk.downcast::<GenericServerKey<AP>>().unwrap()) // We know from the TypeId that

are there other locations where we can use downcast directly ? still feels like there is a copy here potentially from heap to stack e.g. but I don't think there is a way around it

[nsarlin-zama]

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

---

tfhe/src/shortint/backward_compatibility/client_key/mod.rs line 55 at r4 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

standard atomic pattern ?

fixed

---

tfhe/src/shortint/backward_compatibility/server_key/mod.rs line 115 at r4 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

are there other locations where we can use downcast directly ? still feels like there is a copy here potentially from heap to stack e.g. but I don't think there is a way around it

There is a copy but this is not the content of the vec this time, only the pointers and metadata.

---

tfhe/src/shortint/client_key/atomic_pattern/mod.rs line 41 at r3 (raw file):

Previously, IceTDrinker (Arthur Meyre) wrote…

of client keys

done

[nsarlin-zama]

Reviewable status: 62 of 65 files reviewed, 2 unresolved discussions (waiting on @IceTDrinker, @mayeul-zama, and @tmontaigu)

---

tfhe/src/shortint/backward_compatibility/server_key/mod.rs line 115 at r4 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

There is a copy but this is not the content of the vec this time, only the pointers and metadata.

I reused this pattern for the modswitch noise reduction key upgrade

[IceTDrinker]

Thanks !

Reviewed 3 of 3 files at r5, all commit messages.
Reviewable status: all files reviewed, 1 unresolved discussion (waiting on @mayeul-zama, @nsarlin-zama, and @tmontaigu)

[nsarlin-zama]

fixed a conflict

Reviewable status: all files reviewed, 1 unresolved discussion (waiting on @mayeul-zama and @tmontaigu)

[IceTDrinker]

Reviewed 3 of 3 files at r6, all commit messages.
Reviewable status: complete! all files reviewed, all discussions resolved (waiting on @mayeul-zama and @tmontaigu)

[IceTDrinker]

Reviewed 1 of 1 files at r7, all commit messages.
Reviewable status: complete! all files reviewed, all discussions resolved (waiting on @IceTDrinker, @mayeul-zama, and @tmontaigu)