feat(gpu): add support to GPU-accelerated expand to HL's CompactCiphe…

Makefile

@@ -839,7 +839,7 @@ test_high_level_api: install_rs_build_toolchain
 
 test_high_level_api_gpu: install_rs_build_toolchain install_cargo_nextest
 	RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_BUILD_TOOLCHAIN) nextest run --cargo-profile $(CARGO_PROFILE) \
-		--features=integer,internal-keycache,gpu -p $(TFHE_SPEC) \
+		--features=integer,internal-keycache,gpu,zk-pok -p $(TFHE_SPEC) \
 		-E "test(/high_level_api::.*gpu.*/)"
 
 .PHONY: test_strings # Run the tests for strings ci

tfhe/docs/configuration/gpu_acceleration/zk-pok.md

@@ -0,0 +1,81 @@
+# Zero-knowledge proofs
+
+Zero-knowledge proofs (ZK) are a powerful tool to assert that the encryption of a message is correct, as discussed in [advanced features](../../fhe-computation/advanced-features/zk-pok.md).
+However, computation is not possible on the type of ciphertexts it produces (i.e. `ProvenCompactCiphertextList`). This document explains how to use the GPU to accelerate the
+preprocessing step needed to convert ciphertexts formatted for ZK to ciphertexts in the right format for computation purposes on GPU. This 
+operation is called "expansion".
+
+## Proven compact ciphertext list
+
+A proven compact list of ciphertexts can be seen as a compacted collection of ciphertexts which encryption can be verified.
+This verification is currently only supported on the CPU, but the expansion can be accelerated using the GPU.
+This way, verification and expansion can be performed in parallel, efficiently using all the available computational resources.
+
+## Supported types
+Encrypted messages can be integers (like FheUint64) or booleans. The GPU backend does not currently support encrypted strings.
+
+{% hint style="info" %}
+You can enable this feature using the flag: `--features=zk-pok,gpu` when building **TFHE-rs**.
+{% endhint %}
+
+
+## Example
+
+The following example shows how a client can encrypt and prove a ciphertext, and how a server can verify the proof, preprocess the ciphertext and run a computation on it on GPU:
+
+```rust
+use rand::random;
+use tfhe::CompressedServerKey;
+use tfhe::prelude::*;
+use tfhe::set_server_key;
+use tfhe::zk::{CompactPkeCrs, ZkComputeLoad};
+
+pub fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let params = tfhe::shortint::parameters::PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+    // Indicate which parameters to use for the Compact Public Key encryption
+    let cpk_params = tfhe::shortint::parameters::PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+    // And parameters allowing to keyswitch/cast to the computation parameters.
+    let casting_params = tfhe::shortint::parameters::PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+    // Enable the dedicated parameters on the config
+    let config = tfhe::ConfigBuilder::with_custom_parameters(params)
+        .use_dedicated_compact_public_key_parameters((cpk_params, casting_params)).build();
+
+    // The CRS should be generated in an offline phase then shared to all clients and the server
+    let crs = CompactPkeCrs::from_config(config, 64).unwrap();
+
+    // Then use TFHE-rs as usual
+    let client_key = tfhe::ClientKey::generate(config);
+    let compressed_server_key = CompressedServerKey::new(&client_key);
+    let gpu_server_key = compressed_server_key.decompress_to_gpu();
+    let public_key = tfhe::CompactPublicKey::try_new(&client_key).unwrap();
+    // This can be left empty, but if provided allows to tie the proof to arbitrary data
+    let metadata = [b'T', b'F', b'H', b'E', b'-', b'r', b's'];
+
+    let clear_a = random::<u64>();
+    let clear_b = random::<u64>();
+
+    let proven_compact_list = tfhe::ProvenCompactCiphertextList::builder(&public_key)
+        .push(clear_a)
+        .push(clear_b)
+        .build_with_proof_packed(&crs, &metadata, ZkComputeLoad::Verify)?;
+
+    // Server side
+    let result = {
+        set_server_key(gpu_server_key);
+
+        // Verify the ciphertexts
+        let expander =
+            proven_compact_list.verify_and_expand(&crs, &public_key, &metadata)?;
+        let a: tfhe::FheUint64 = expander.get(0)?.unwrap();
+        let b: tfhe::FheUint64 = expander.get(1)?.unwrap();
+
+        a + b
+    };
+
+    // Back on the client side
+    let a_plus_b: u64 = result.decrypt(&client_key);
+    assert_eq!(a_plus_b, clear_a.wrapping_add(clear_b));
+
+    Ok(())
+}
+```

tfhe/src/core_crypto/gpu/entities/lwe_ciphertext_list.rs

@@ -7,7 +7,7 @@ use crate::core_crypto::prelude::{
 use tfhe_cuda_backend::cuda_bind::cuda_memcpy_async_gpu_to_gpu;
 
 /// A structure representing a vector of LWE ciphertexts with 64 bits of precision on the GPU.
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct CudaLweCiphertextList<T: UnsignedInteger>(pub(crate) CudaLweList<T>);
 
 #[allow(dead_code)]

tfhe/src/core_crypto/gpu/entities/lwe_compact_ciphertext_list.rs

@@ -4,12 +4,33 @@
 use crate::core_crypto::gpu::vec::CudaVec;
 use crate::core_crypto::gpu::{CudaLweList, CudaStreams};
 use crate::core_crypto::prelude::{
-    lwe_compact_ciphertext_list_size, Container, LweCompactCiphertextList, UnsignedInteger,
+    lwe_compact_ciphertext_list_size, CiphertextModulus, Container, LweCiphertextCount,
+    LweCompactCiphertextList, LweDimension, UnsignedInteger,
 };
 
 pub struct CudaLweCompactCiphertextList<T: UnsignedInteger>(pub CudaLweList<T>);
 
 impl<T: UnsignedInteger> CudaLweCompactCiphertextList<T> {
+    pub fn duplicate(&self, streams: &CudaStreams) -> Self {
+        Self(self.0.duplicate(streams))
+    }
+
+    pub fn from_d_vec(
+        d_vec: CudaVec<T>,
+        lwe_ciphertext_count: LweCiphertextCount,
+        lwe_dimension: LweDimension,
+        ciphertext_modulus: CiphertextModulus<T>,
+    ) -> Self {
+        Self {
+            0: CudaLweList {
+                d_vec,
+                lwe_ciphertext_count,
+                lwe_dimension,
+                ciphertext_modulus,
+            },
+        }
+    }
+
     pub fn from_lwe_compact_ciphertext_list<C: Container<Element = T>>(
         h_ct: &LweCompactCiphertextList<C>,
         streams: &CudaStreams,

tfhe/src/core_crypto/gpu/entities/lwe_keyswitch_key.rs

@@ -10,6 +10,7 @@ use crate::core_crypto::prelude::{
     UnsignedInteger,
 };
 
+#[derive(Clone)]
 #[allow(dead_code)]
 pub struct CudaLweKeyswitchKey<T: UnsignedInteger> {
     pub(crate) d_vec: CudaVec<T>,

tfhe/src/core_crypto/gpu/mod.rs

@@ -914,7 +914,7 @@ pub unsafe fn fourier_transform_backward_as_torus_f128_async<T: UnsignedInteger>
     );
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct CudaLweList<T: UnsignedInteger> {
     // Pointer to GPU data
     pub d_vec: CudaVec<T>,
@@ -926,6 +926,17 @@ pub struct CudaLweList<T: UnsignedInteger> {
     pub ciphertext_modulus: CiphertextModulus<T>,
 }
 
+impl<T: UnsignedInteger> CudaLweList<T> {
+    pub fn duplicate(&self, streams: &CudaStreams) -> Self {
+        Self {
+            d_vec: self.d_vec.duplicate(streams),
+            lwe_ciphertext_count: self.lwe_ciphertext_count,
+            lwe_dimension: self.lwe_dimension,
+            ciphertext_modulus: self.ciphertext_modulus,
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct CudaGlweList<T: UnsignedInteger> {
     // Pointer to GPU data

tfhe/src/high_level_api/backward_compatibility/compact_list.rs

@@ -3,6 +3,9 @@ use tfhe_versionable::{Upgrade, Version, VersionsDispatch};
 
 use crate::{CompactCiphertextList, Tag};
 
+#[cfg(feature = "zk-pok")]
+use crate::ProvenCompactCiphertextList;
+
 #[derive(Version)]
 pub struct CompactCiphertextListV0(crate::integer::ciphertext::CompactCiphertextList);
 
@@ -11,15 +14,12 @@ impl Upgrade<CompactCiphertextList> for CompactCiphertextListV0 {
 
     fn upgrade(self) -> Result<CompactCiphertextList, Self::Error> {
         Ok(CompactCiphertextList {
-            inner: self.0,
+            inner: crate::high_level_api::compact_list::InnerCompactCiphertextList::Cpu(self.0),
             tag: Tag::default(),
         })
     }
 }
 
-#[cfg(feature = "zk-pok")]
-use crate::ProvenCompactCiphertextList;
-
 #[derive(VersionsDispatch)]
 pub enum CompactCiphertextListVersions {
     V0(CompactCiphertextListV0),