Skip to content

Commit

Permalink
refactor(integer): add compression key types
Browse files Browse the repository at this point in the history
  • Loading branch information
mayeul-zama committed Sep 11, 2024
1 parent 93ff699 commit c0d9839
Show file tree
Hide file tree
Showing 14 changed files with 215 additions and 45 deletions.
63 changes: 56 additions & 7 deletions tfhe/src/high_level_api/backward_compatibility/keys.rs
Original file line number Diff line number Diff line change
Expand Up @@ -211,14 +211,35 @@ impl Upgrade<IntegerClientKeyV1> for IntegerClientKeyV0 {
}
}

impl Upgrade<IntegerClientKey> for IntegerClientKeyV1 {
impl Upgrade<IntegerClientKeyV2> for IntegerClientKeyV1 {
type Error = Infallible;

fn upgrade(self) -> Result<IntegerClientKeyV2, Self::Error> {
Ok(IntegerClientKeyV2 {
key: self.key,
dedicated_compact_private_key: self.dedicated_compact_private_key,
compression_key: self.compression_key,
})
}
}

#[derive(Version)]
pub(crate) struct IntegerClientKeyV2 {
pub(crate) key: crate::integer::ClientKey,
pub(crate) dedicated_compact_private_key: Option<CompactPrivateKey>,
pub(crate) compression_key: Option<crate::shortint::list_compression::CompressionPrivateKeys>,
}

impl Upgrade<IntegerClientKey> for IntegerClientKeyV2 {
type Error = Infallible;

fn upgrade(self) -> Result<IntegerClientKey, Self::Error> {
Ok(IntegerClientKey {
key: self.key,
dedicated_compact_private_key: self.dedicated_compact_private_key,
compression_key: self.compression_key,
compression_key: self
.compression_key
.map(|key| crate::integer::compression_keys::CompressionPrivateKeys { key }),
})
}
}
Expand All @@ -228,7 +249,8 @@ impl Upgrade<IntegerClientKey> for IntegerClientKeyV1 {
pub(crate) enum IntegerClientKeyVersions {
V0(IntegerClientKeyV0),
V1(IntegerClientKeyV1),
V2(IntegerClientKey),
V2(IntegerClientKeyV2),
V3(IntegerClientKey),
}

#[derive(Version)]
Expand Down Expand Up @@ -261,11 +283,11 @@ impl Upgrade<IntegerServerKeyV1> for IntegerServerKeyV0 {
}
}

impl Upgrade<IntegerServerKey> for IntegerServerKeyV1 {
impl Upgrade<IntegerServerKeyV2> for IntegerServerKeyV1 {
type Error = Infallible;

fn upgrade(self) -> Result<IntegerServerKey, Self::Error> {
Ok(IntegerServerKey {
fn upgrade(self) -> Result<IntegerServerKeyV2, Self::Error> {
Ok(IntegerServerKeyV2 {
key: self.key,
cpk_key_switching_key_material: self.cpk_key_switching_key_material,
compression_key: self.compression_key,
Expand All @@ -274,11 +296,38 @@ impl Upgrade<IntegerServerKey> for IntegerServerKeyV1 {
}
}

#[derive(Version)]
pub struct IntegerServerKeyV2 {
pub(crate) key: crate::integer::ServerKey,
pub(crate) cpk_key_switching_key_material:
Option<crate::integer::key_switching_key::KeySwitchingKeyMaterial>,
pub(crate) compression_key: Option<crate::shortint::list_compression::CompressionKey>,
pub(crate) decompression_key: Option<crate::shortint::list_compression::DecompressionKey>,
}

impl Upgrade<IntegerServerKey> for IntegerServerKeyV2 {
type Error = Infallible;

fn upgrade(self) -> Result<IntegerServerKey, Self::Error> {
Ok(IntegerServerKey {
key: self.key,
cpk_key_switching_key_material: self.cpk_key_switching_key_material,
compression_key: self
.compression_key
.map(|key| crate::integer::compression_keys::CompressionKey { key }),
decompression_key: self
.decompression_key
.map(|key| crate::integer::compression_keys::DecompressionKey { key }),
})
}
}

#[derive(VersionsDispatch)]
pub enum IntegerServerKeyVersions {
V0(IntegerServerKeyV0),
V1(IntegerServerKeyV1),
V2(IntegerServerKey),
V2(IntegerServerKeyV2),
V3(IntegerServerKey),
}

#[derive(Version)]
Expand Down
2 changes: 1 addition & 1 deletion tfhe/src/high_level_api/keys/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@ use super::{CompressedServerKey, ServerKey};
use crate::high_level_api::backward_compatibility::keys::ClientKeyVersions;
use crate::high_level_api::config::Config;
use crate::high_level_api::keys::{CompactPrivateKey, IntegerClientKey};
use crate::integer::compression_keys::CompressionPrivateKeys;
use crate::prelude::Tagged;
use crate::shortint::list_compression::CompressionPrivateKeys;
use crate::shortint::MessageModulus;
use crate::Tag;
use concrete_csprng::seeders::Seed;
Expand Down
17 changes: 8 additions & 9 deletions tfhe/src/high_level_api/keys/inner.rs
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
use crate::core_crypto::commons::generators::DeterministicSeeder;
use crate::core_crypto::prelude::ActivatedRandomGenerator;
use crate::high_level_api::backward_compatibility::keys::*;
use crate::integer::public_key::CompactPublicKey;
use crate::integer::CompressedCompactPublicKey;
use crate::shortint::list_compression::{
use crate::integer::compression_keys::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, CompressionPrivateKeys,
DecompressionKey,
};
use crate::integer::public_key::CompactPublicKey;
use crate::integer::CompressedCompactPublicKey;
use crate::shortint::parameters::list_compression::CompressionParameters;
use crate::shortint::MessageModulus;
use crate::Error;
Expand Down Expand Up @@ -103,11 +103,11 @@ impl IntegerClientKey {
let cks = crate::shortint::engine::ShortintEngine::new_from_seeder(&mut seeder)
.new_client_key(config.block_parameters.into());

let key = crate::integer::ClientKey::from(cks);

let compression_key = config
.compression_parameters
.map(|params| cks.new_compression_private_key(params));

let key = crate::integer::ClientKey::from(cks);
.map(|params| key.new_compression_private_key(params));

let dedicated_compact_private_key = config
.dedicated_compact_public_key_parameters
Expand Down Expand Up @@ -193,7 +193,7 @@ impl From<IntegerConfig> for IntegerClientKey {

let compression_key = config
.compression_parameters
.map(|params| key.key.new_compression_private_key(params));
.map(|params| key.new_compression_private_key(params));

Self {
key,
Expand Down Expand Up @@ -225,7 +225,7 @@ impl IntegerServerKey {
|| (None, None),
|a| {
let (compression_key, decompression_key) =
cks.key.new_compression_decompression_keys(a);
cks.new_compression_decompression_keys(a);
(Some(compression_key), Some(decompression_key))
},
);
Expand Down Expand Up @@ -312,7 +312,6 @@ impl IntegerCompressedServerKey {
.as_ref()
.map_or((None, None), |compression_private_key| {
let (compression_keys, decompression_keys) = client_key
.key
.key
.new_compressed_compression_decompression_keys(compression_private_key);

Expand Down
4 changes: 2 additions & 2 deletions tfhe/src/high_level_api/keys/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@ use crate::backward_compatibility::keys::{CompressedServerKeyVersions, ServerKey
#[cfg(feature = "gpu")]
use crate::core_crypto::gpu::{synchronize_devices, CudaStreams};
use crate::high_level_api::keys::{IntegerCompressedServerKey, IntegerServerKey};
use crate::prelude::Tagged;
use crate::shortint::list_compression::{
use crate::integer::compression_keys::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, DecompressionKey,
};
use crate::prelude::Tagged;
use crate::shortint::MessageModulus;
use crate::Tag;
use std::sync::Arc;
Expand Down
30 changes: 30 additions & 0 deletions tfhe/src/integer/backward_compatibility/list_compression.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
use crate::integer::compression_keys::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, CompressionPrivateKeys,
DecompressionKey,
};
use tfhe_versionable::VersionsDispatch;

#[derive(VersionsDispatch)]
pub enum CompressionKeyVersions {
V0(CompressionKey),
}

#[derive(VersionsDispatch)]
pub enum DecompressionKeyVersions {
V0(DecompressionKey),
}

#[derive(VersionsDispatch)]
pub enum CompressedCompressionKeyVersions {
V0(CompressedCompressionKey),
}

#[derive(VersionsDispatch)]
pub enum CompressedDecompressionKeyVersions {
V0(CompressedDecompressionKey),
}

#[derive(VersionsDispatch)]
pub enum CompressionPrivateKeysVersions {
V0(CompressionPrivateKeys),
}
1 change: 1 addition & 0 deletions tfhe/src/integer/backward_compatibility/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
pub mod ciphertext;
pub mod client_key;
pub mod key_switching_key;
pub mod list_compression;
pub mod public_key;
pub mod server_key;
pub mod wopbs;
8 changes: 5 additions & 3 deletions tfhe/src/integer/ciphertext/compressed_ciphertext_list.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use super::{DataKind, Expandable, RadixCiphertext, SignedRadixCiphertext};
use crate::integer::backward_compatibility::ciphertext::CompressedCiphertextListVersions;
use crate::integer::compression_keys::{CompressionKey, DecompressionKey};
use crate::integer::BooleanBlock;
use crate::shortint::ciphertext::CompressedCiphertextList as ShortintCompressedCiphertextList;
use crate::shortint::list_compression::{CompressionKey, DecompressionKey};
use crate::shortint::Ciphertext;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
Expand Down Expand Up @@ -84,7 +84,9 @@ impl CompressedCiphertextListBuilder {
}

pub fn build(&self, comp_key: &CompressionKey) -> CompressedCiphertextList {
let packed_list = comp_key.compress_ciphertexts_into_list(&self.ciphertexts);
let packed_list = comp_key
.key
.compress_ciphertexts_into_list(&self.ciphertexts);

CompressedCiphertextList {
packed_list,
Expand Down Expand Up @@ -128,7 +130,7 @@ impl CompressedCiphertextList {
Some((
(start_block_index..end_block_index)
.into_par_iter()
.map(|i| decomp_key.unpack(&self.packed_list, i).unwrap())
.map(|i| decomp_key.key.unpack(&self.packed_list, i).unwrap())
.collect(),
current_info,
))
Expand Down
16 changes: 12 additions & 4 deletions tfhe/src/integer/client_key/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@ use crate::integer::block_decomposition::BlockRecomposer;
use crate::integer::ciphertext::boolean_value::BooleanBlock;
use crate::integer::ciphertext::{CompressedCrtCiphertext, CrtCiphertext};
use crate::integer::client_key::utils::i_crt;
use crate::integer::compression_keys::{CompressionKey, CompressionPrivateKeys, DecompressionKey};
use crate::integer::encryption::{encrypt_crt, encrypt_words_radix_impl};
use crate::shortint::ciphertext::Degree;
use crate::shortint::list_compression::{CompressionKey, CompressionPrivateKeys, DecompressionKey};
use crate::shortint::parameters::{CompressionParameters, MessageModulus};
use crate::shortint::{
Ciphertext, ClientKey as ShortintClientKey, ShortintParameterSet as ShortintParameters,
Expand Down Expand Up @@ -720,14 +720,22 @@ impl ClientKey {
&self,
params: CompressionParameters,
) -> CompressionPrivateKeys {
self.key.new_compression_private_key(params)
CompressionPrivateKeys {
key: self.key.new_compression_private_key(params),
}
}

pub fn new_compression_decompression_keys(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> (CompressionKey, DecompressionKey) {
self.key
.new_compression_decompression_keys(private_compression_key)
let (comp_key, decomp_key) = self
.key
.new_compression_decompression_keys(&private_compression_key.key);

(
CompressionKey { key: comp_key },
DecompressionKey { key: decomp_key },
)
}
}
14 changes: 10 additions & 4 deletions tfhe/src/integer/client_key/radix.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@ use crate::core_crypto::prelude::{SignedNumeric, UnsignedNumeric};
use crate::integer::backward_compatibility::client_key::RadixClientKeyVersions;
use crate::integer::block_decomposition::{DecomposableInto, RecomposableFrom};
use crate::integer::ciphertext::{RadixCiphertext, SignedRadixCiphertext};
use crate::integer::BooleanBlock;
use crate::shortint::list_compression::{
use crate::integer::compression_keys::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionPrivateKeys,
};
use crate::integer::BooleanBlock;
use crate::shortint::{Ciphertext as ShortintCiphertext, PBSParameters as ShortintParameters};
use serde::{Deserialize, Serialize};
use tfhe_versionable::Versionize;
Expand Down Expand Up @@ -139,9 +139,15 @@ impl RadixClientKey {
&self,
private_compression_key: &CompressionPrivateKeys,
) -> (CompressedCompressionKey, CompressedDecompressionKey) {
self.key
let (comp_key, decomp_key) = self
.key
.new_compressed_compression_decompression_keys(private_compression_key)
.key
.new_compressed_compression_decompression_keys(&private_compression_key.key);

(
CompressedCompressionKey { key: comp_key },
CompressedDecompressionKey { key: decomp_key },
)
}
}

Expand Down
66 changes: 66 additions & 0 deletions tfhe/src/integer/compression_keys.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
use super::ClientKey;
use crate::integer::backward_compatibility::list_compression::*;
use serde::{Deserialize, Serialize};
use tfhe_versionable::Versionize;

#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
#[versionize(CompressionPrivateKeysVersions)]
pub struct CompressionPrivateKeys {
pub(crate) key: crate::shortint::list_compression::CompressionPrivateKeys,
}

#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
#[versionize(CompressionKeyVersions)]
pub struct CompressionKey {
pub(crate) key: crate::shortint::list_compression::CompressionKey,
}

#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
#[versionize(DecompressionKeyVersions)]
pub struct DecompressionKey {
pub(crate) key: crate::shortint::list_compression::DecompressionKey,
}

#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
#[versionize(CompressedCompressionKeyVersions)]
pub struct CompressedCompressionKey {
pub(crate) key: crate::shortint::list_compression::CompressedCompressionKey,
}

#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
#[versionize(CompressedDecompressionKeyVersions)]
pub struct CompressedDecompressionKey {
pub(crate) key: crate::shortint::list_compression::CompressedDecompressionKey,
}

impl CompressedCompressionKey {
pub fn decompress(&self) -> CompressionKey {
CompressionKey {
key: self.key.decompress(),
}
}
}

impl CompressedDecompressionKey {
pub fn decompress(&self) -> DecompressionKey {
DecompressionKey {
key: self.key.decompress(),
}
}
}

impl ClientKey {
pub fn new_compressed_compression_decompression_keys(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> (CompressedCompressionKey, CompressedDecompressionKey) {
let (comp_key, decomp_key) = self
.key
.new_compressed_compression_decompression_keys(&private_compression_key.key);

(
CompressedCompressionKey { key: comp_key },
CompressedDecompressionKey { key: decomp_key },
)
}
}
Loading

0 comments on commit c0d9839

Please sign in to comment.