Ns/feat/atomic patterns #2024

nsarlin-zama · 2025-01-31T17:47:00Z

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 ?

tmontaigu

Reviewed 8 of 45 files at r1, 5 of 14 files at r2, 51 of 74 files at r3, 36 of 56 files at r4, 179 of 179 files at r5, all commit messages.
Reviewable status: all files reviewed, 8 unresolved discussions (waiting on @mayeul-zama and @nsarlin-zama)

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

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

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

I don't have one

tfhe/src/core_crypto/entities/compressed_modulus_switched_lwe_ciphertext.rs line 78 at r3 (raw file):

        log_modulus: CiphertextModulusLog,
    ) -> Self
    where

minor suggestion:
we could have kept the Cont as the only generic and have

where 
Cont: Container,
Cont::Element: UnsignedInteger + CastInto<PackingScalar>

Then in the assert Cont::Element::BITS should work or the more verbose <Cont::Element as UnsignedInteger>::BITS

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

    impl<
            AP: 'static

I tend to think that a where clause would look better

tfhe/src/shortint/parameters/mod.rs line 464 at r5 (raw file):

            ShortintParameterSetInner::KS32PBS(params) => {
                let distribution = params.lwe_noise_distribution();
                match distribution {

Isn't there a From that is implemented ? I think I saw one 🤔

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

    // This buffer will be converted when needed into temporary lwe ciphertexts, eventually by
    // splitting u128 blocks into smaller scalars
    buffer: Vec<u128>,

Since this now acts as a buffer that not always returns slices of u64, should this be a Vec ?
Was u128 chosen because it would give some alignment properties ?

nsarlin-zama

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

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

Previously, tmontaigu (tmontaigu) wrote…

Since this now acts as a buffer that not always returns slices of u64, should this be a Vec ?
Was u128 chosen because it would give some alignment properties ?

Yes, the idea is to allocate with the strictest alignment so that any smaller scalar will also work.
What do you suggest instead of a vec? There might be a cleaner/simpler way that I missed!

tfhe/src/shortint/parameters/mod.rs line 464 at r5 (raw file):

Previously, tmontaigu (tmontaigu) wrote…

Isn't there a From that is implemented ? I think I saw one 🤔

yes you're right, I forgot about it

nsarlin-zama

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

tfhe/src/core_crypto/entities/compressed_modulus_switched_lwe_ciphertext.rs line 78 at r3 (raw file):

Previously, tmontaigu (tmontaigu) wrote…

minor suggestion:
we could have kept the Cont as the only generic and have
where 
Cont: Container,
Cont::Element: UnsignedInteger + CastInto<PackingScalar>
Then in the assert Cont::Element::BITS should work or the more verbose <Cont::Element as UnsignedInteger>::BITS

done

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

Previously, tmontaigu (tmontaigu) wrote…

I tend to think that a where clause would look better

done

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

Previously, tmontaigu (tmontaigu) wrote…

I don't have one

I renamed it to simply sk.owned()

nsarlin-zama

Reviewable status: 30 of 230 files reviewed, 4 unresolved discussions (waiting on @mayeul-zama and @tmontaigu)

tfhe/src/shortint/parameters/mod.rs line 464 at r5 (raw file):

Previously, nsarlin-zama (Nicolas Sarlin) wrote…

yes you're right, I forgot about it

fixed

Since only one kind is used at a time we don't need do allocate both

This allows to call it without having to define a max_noise_level

mayeul-zama · 2025-04-11T09:53:49Z

tfhe/src/shortint/server_key/mod.rs

@@ -430,6 +430,30 @@ pub struct ServerKey {
    pub pbs_order: PBSOrder,
 }

+/// Represents the number of elements in a [`LookupTable`] represented by a Glwe ciphertext
+#[derive(Copy, Clone, Debug)]
+pub struct LookupTableSize {


It can be used for any GLWE, not only LUTs
So it could be called GlweParameters

Not sure having this type is very useful

This is supposed to be a bit more high level that the Glwe. Lookup table is at the shortint level and Glwe is at the core level. In theory I think you could use shortint without knowing what a glwe is.

The main goal is to reduce the number of parameters in the lut creation functions, because these two always go together.
And in the AtomicPattern, since we always use them together we can do let size = ap.lookup_table_size() in one call instead of having to do:

let glwe_size = ap.glwe_size(); let polynomial_size = ap.polynomial_size();

This is not crucial however!

mayeul-zama · 2025-04-11T12:30:01Z

tfhe/src/shortint/atomic_pattern/classical.rs

Modifications to this file in commit "insert the AP inside the ServerKey" could be part of commit "create atomic pattern trait and enum"

mayeul-zama · 2025-04-11T12:30:57Z

tfhe/src/shortint/atomic_pattern/mod.rs

Modifications to this file in commit "insert the AP inside the ServerKey" could be part of commit "create atomic pattern trait and enum"

mayeul-zama · 2025-04-11T12:39:39Z

tfhe/src/shortint/atomic_pattern/mod.rs

@@ -48,7 +50,7 @@ pub enum AtomicPatternKind {
 /// The atomic pattern can be seen as a black box that will apply a lookup table and refresh the
 /// ciphertext noise to a nominal level. Between applications of the AP, it is possible to do a
 /// certain number of linear operations.
-pub trait AtomicPattern {
+pub trait AtomicPattern: std::fmt::Debug {


Why is Debug necessary?

Without I cannot derive Debug for AtomicPatternServerKey, and this is needed for the ServerKey. I can certainly move the bound to the DynamicAtomicPattern trait if you prefer

I moved it to the DAP trait

mayeul-zama · 2025-04-11T12:40:56Z

tfhe/src/shortint/engine/mod.rs

        };

-        LweCiphertextMutView::from_container(buffer, ciphertext_modulus)
+        // This should not panic as long as `Scalar::BITS` is a divisor of 128


Maybe that could be asserted in this function

I added a panic with a more explicit error message

mayeul-zama · 2025-04-11T12:50:03Z

tfhe/src/core_crypto/entities/lwe_keyswitch_key.rs

@@ -455,7 +457,11 @@ impl<C: Container<Element = u64>> ParameterSetConformant for LweKeyswitchKey<C>
            ciphertext_modulus,
        } = self;

-        *ciphertext_modulus == parameter_set.ciphertext_modulus
+        let Ok(parameters_modulus) = parameter_set.ciphertext_modulus.try_to() else {


What is this try_to?

This is a method on the CiphertextModulus:

tfhe-rs/tfhe/src/core_crypto/commons/ciphertext_modulus.rs

Line 302 in 1b5f52c

pub fn try_to<ScalarTo: UnsignedInteger + CastInto<u128>>(

Because in the parameters the CiphertextModulus is u64 and in LweKeyswitchKey it is a generic, so this checks that they are compatible

mayeul-zama · 2025-04-11T12:52:41Z

tfhe/src/shortint/parameters/mod.rs

@@ -433,6 +459,9 @@ impl ShortintParameterSet {
            ShortintParameterSetInner::PBSOnly(params) => params.lwe_noise_distribution(),
            ShortintParameterSetInner::WopbsOnly(params) => params.lwe_noise_distribution,
            ShortintParameterSetInner::PBSAndWopbs(params, _) => params.lwe_noise_distribution(),
+            ShortintParameterSetInner::KS32PBS(params) => {
+                params.lwe_noise_distribution().to_u64_distribution()


Maybe this function should not return a hardcoded DynamicDistribution<u64>

What do you have in mind ? I can't change the return type dynamically.
But imo this is not a problem because the generic is only used in the TUniform, and a u32 bound is also a u64 bound so all DynamicDistribution<u32> are valid DynamicDistribution<u64>

cla-bot bot added the cla-signed label Jan 31, 2025

tmontaigu self-requested a review February 5, 2025 09:22

tmontaigu reviewed Feb 5, 2025

View reviewed changes

nsarlin-zama commented Feb 5, 2025

View reviewed changes

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch from 41c1d18 to a413178 Compare February 5, 2025 14:53

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch 5 times, most recently from b11e47a to 1fc1441 Compare March 19, 2025 17:38

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch 11 times, most recently from 14a920c to 600a1b7 Compare March 27, 2025 09:51

nsarlin-zama mentioned this pull request Mar 27, 2025

Ns/chore/conformance unwrap #2212

Merged

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch 3 times, most recently from 20843dd to 4a759f6 Compare April 4, 2025 16:09

nsarlin-zama changed the base branch from main to ns/feat/versionize_skip April 4, 2025 16:10

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch from 4a759f6 to 74d237b Compare April 4, 2025 16:29

Base automatically changed from ns/feat/versionize_skip to main April 7, 2025 07:48

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch from 74d237b to 7707a5b Compare April 7, 2025 08:47

nsarlin-zama marked this pull request as ready for review April 7, 2025 08:59

tmontaigu reviewed Apr 11, 2025

View reviewed changes

nsarlin-zama commented Apr 11, 2025

View reviewed changes

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch from d9949dd to 6336cfc Compare April 11, 2025 09:56

nsarlin-zama commented Apr 11, 2025

View reviewed changes

nsarlin-zama added 6 commits April 11, 2025 12:08

refactor(shortint): wrap PbsOrder into AtomicPattern in ciphertext

aec46e0

refactor(shortint): factorize generate_lookup_table

dd9669c

refactor(shortint): use a single lwe buffer inside shortint engine

f1e4857

Since only one kind is used at a time we don't need do allocate both

refactor(shortint): function to directly set noise level to nominal

c81cb11

This allows to call it without having to define a max_noise_level

refactor(shortint): use a dedicated type for lut size

313937f

refactor(shortint): remove degree in generate_lookup_table_no_encode

efabfbe

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch from 6336cfc to f2161a4 Compare April 11, 2025 10:13

mayeul-zama reviewed Apr 11, 2025

View reviewed changes

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch from f2161a4 to 3693454 Compare April 11, 2025 16:13

nsarlin-zama added 13 commits April 11, 2025 18:18

feat(shortint): create atomic pattern trait and enum

624bf64

feat(shortint): insert the AP inside the ServerKey

b334c01

refactor(shortint): engine can create any atomic pattern sk

5bc0258

feat(shortint): add the dynamic ap

f5ec05d

refactor(shortint): support any ciphertext modulus in the engine

574cc71

refactor(core): allow different input/output scalars in multibit br

237ce11

refactor(shortint): support any scalar in modswitch noise reduction

335de23

refactor(shortint): make oprf generic over the Scalar type

d218e63

refactor(shortint): make modswitch compression generic over scalar

033e57d

refactor(core): make ksk generation generic over the scalar type

2090e6b

feat(shortint): introduce the KS32 atomic pattern

feac22a

chore(tests): add support for AP in tests and benches

b0db18e

chore(shortint): add tests for the KS32 AP

fc43be6

nsarlin-zama force-pushed the ns/feat/atomic_patterns branch from 3693454 to fc43be6 Compare April 11, 2025 16:19

nsarlin-zama requested review from tmontaigu and mayeul-zama April 11, 2025 16:20

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Ns/feat/atomic patterns #2024

Ns/feat/atomic patterns #2024

nsarlin-zama commented Jan 31, 2025 •

edited

Loading

tmontaigu left a comment

nsarlin-zama left a comment

tmontaigu left a comment

nsarlin-zama left a comment

nsarlin-zama left a comment

nsarlin-zama left a comment

mayeul-zama Apr 11, 2025

nsarlin-zama Apr 11, 2025

mayeul-zama Apr 11, 2025

nsarlin-zama Apr 11, 2025

mayeul-zama Apr 11, 2025

nsarlin-zama Apr 11, 2025

mayeul-zama Apr 11, 2025

nsarlin-zama Apr 11, 2025

nsarlin-zama Apr 11, 2025

mayeul-zama Apr 11, 2025

nsarlin-zama Apr 11, 2025

mayeul-zama Apr 11, 2025

nsarlin-zama Apr 11, 2025

mayeul-zama Apr 11, 2025

nsarlin-zama Apr 11, 2025

Ns/feat/atomic patterns #2024

Are you sure you want to change the base?

Ns/feat/atomic patterns #2024

Conversation

nsarlin-zama commented Jan 31, 2025 • edited Loading

PR content/description

Atomic Pattern

Implementation

Dynamic AP

KS32

Check-list:

tmontaigu left a comment

Choose a reason for hiding this comment

nsarlin-zama left a comment

Choose a reason for hiding this comment

tmontaigu left a comment

Choose a reason for hiding this comment

nsarlin-zama left a comment

Choose a reason for hiding this comment

nsarlin-zama left a comment

Choose a reason for hiding this comment

nsarlin-zama left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nsarlin-zama commented Jan 31, 2025 •

edited

Loading