From 92163c2646047f9320131339bff9575640ef79a5 Mon Sep 17 00:00:00 2001 From: tmontaigu Date: Wed, 11 Dec 2024 11:48:31 +0100 Subject: [PATCH] chore(hlapi): Add array conversion from/to Vec Add `From` impl to allow conversion from Vec like Vec to Cpu/Gpu array. --- tfhe/src/high_level_api/array/cpu/booleans.rs | 24 ++- tfhe/src/high_level_api/array/cpu/integers.rs | 44 +++++- .../high_level_api/array/dynamic/booleans.rs | 39 ++++- tfhe/src/high_level_api/array/gpu/booleans.rs | 27 +++- tfhe/src/high_level_api/array/gpu/integers.rs | 50 ++++++- tfhe/src/high_level_api/array/tests/mod.rs | 137 +++++++++++++++++- tfhe/src/high_level_api/booleans/inner.rs | 3 +- 7 files changed, 316 insertions(+), 8 deletions(-) diff --git a/tfhe/src/high_level_api/array/cpu/booleans.rs b/tfhe/src/high_level_api/array/cpu/booleans.rs index d41a17fbcd..fc46c7cfa6 100644 --- a/tfhe/src/high_level_api/array/cpu/booleans.rs +++ b/tfhe/src/high_level_api/array/cpu/booleans.rs @@ -7,7 +7,7 @@ use crate::high_level_api::array::{ArrayBackend, BackendDataContainer, BackendDa use crate::high_level_api::global_state; use crate::integer::BooleanBlock; use crate::prelude::{FheDecrypt, FheTryEncrypt}; -use crate::{ClientKey, FheId}; +use crate::{ClientKey, FheBool, FheId, Tag}; use rayon::prelude::*; use std::ops::RangeBounds; @@ -36,6 +36,28 @@ impl ArrayBackend for CpuFheBoolArrayBackend { type Owned = Vec; } +impl From> for CpuFheBoolArray { + fn from(value: Vec) -> Self { + let vec = value + .into_iter() + .map(|b| BooleanBlock::new_unchecked(b.into_raw_parts())) + .collect::>(); + + let shape = vec![vec.len()]; + Self::new(vec, shape) + } +} + +impl From for Vec { + fn from(value: CpuFheBoolArray) -> Self { + value + .into_container() + .into_iter() + .map(|boolean_block| FheBool::new(boolean_block, Tag::default())) + .collect() + } +} + impl BackendDataContainer for &[BooleanBlock] { type Backend = CpuFheBoolArrayBackend; diff --git a/tfhe/src/high_level_api/array/cpu/integers.rs b/tfhe/src/high_level_api/array/cpu/integers.rs index e478ea87fd..ddf41e942e 100644 --- a/tfhe/src/high_level_api/array/cpu/integers.rs +++ b/tfhe/src/high_level_api/array/cpu/integers.rs @@ -19,7 +19,7 @@ use crate::integer::server_key::radix_parallel::scalar_div_mod::SignedReciprocab use crate::integer::server_key::{Reciprocable, ScalarMultiplier}; use crate::integer::{IntegerRadixCiphertext, RadixCiphertext, SignedRadixCiphertext}; use crate::prelude::{FheDecrypt, FheTryEncrypt}; -use crate::{ClientKey, Error}; +use crate::{ClientKey, Error, FheInt, FheUint, Tag}; use rayon::prelude::*; use std::marker::PhantomData; use std::ops::RangeBounds; @@ -54,6 +54,48 @@ where type Owned = Vec; } +impl From>> for CpuFheUintArray { + fn from(value: Vec>) -> Self { + let vec: Vec<_> = value + .into_iter() + .map(|uint| uint.into_raw_parts().0) + .collect(); + let shape = vec![vec.len()]; + Self::new(vec, shape) + } +} + +impl From> for Vec> { + fn from(value: CpuFheUintArray) -> Self { + value + .into_container() + .into_iter() + .map(|radix| FheUint::new(radix, Tag::default())) + .collect() + } +} + +impl From>> for CpuFheIntArray { + fn from(value: Vec>) -> Self { + let vec: Vec<_> = value + .into_iter() + .map(|uint| uint.into_raw_parts().0) + .collect(); + let shape = vec![vec.len()]; + Self::new(vec, shape) + } +} + +impl From> for Vec> { + fn from(value: CpuFheIntArray) -> Self { + value + .into_container() + .into_iter() + .map(|radix| FheInt::new(radix, Tag::default())) + .collect() + } +} + #[inline] #[track_caller] fn par_map_sks_op_on_pair_of_elements<'a, T, F>( diff --git a/tfhe/src/high_level_api/array/dynamic/booleans.rs b/tfhe/src/high_level_api/array/dynamic/booleans.rs index 3f3b0ade81..b412da7617 100644 --- a/tfhe/src/high_level_api/array/dynamic/booleans.rs +++ b/tfhe/src/high_level_api/array/dynamic/booleans.rs @@ -11,7 +11,7 @@ use super::super::{FheBackendArray, FheBackendArraySlice, FheBackendArraySliceMu use crate::array::traits::TensorSlice; use crate::integer::BooleanBlock; use crate::prelude::{FheDecrypt, FheTryEncrypt}; -use crate::{ClientKey, Device}; +use crate::{ClientKey, CpuFheBoolArray, Device, FheBool}; use std::borrow::{Borrow, Cow}; use std::ops::RangeBounds; @@ -33,6 +33,43 @@ impl ArrayBackend for DynFheBoolArrayBackend { type Owned = InnerBoolArray; } +impl TryFrom> for FheBoolArray { + type Error = crate::Error; + + fn try_from(values: Vec) -> Result { + if values.is_empty() { + return Ok(Self::new(InnerBoolArray::Cpu(vec![]), vec![0])); + } + + let shape = vec![values.len()]; + let device_of_first = values[0].current_device(); + let inner = match device_of_first { + Device::Cpu => { + let new_values = values + .into_iter() + .map(|value| value.ciphertext.into_cpu()) + .collect::>(); + + InnerBoolArray::Cpu(new_values) + } + #[cfg(feature = "gpu")] + Device::CudaGpu => return Err(crate::error!("Array do not support GPU")), + #[cfg(feature = "hpu")] + Device::Hpu => return Err(crate::error!("Array do not support HPU")), + }; + + Ok(Self::new(inner, shape)) + } +} + +impl From for FheBoolArray { + fn from(cpu_array: CpuFheBoolArray) -> Self { + let CpuFheBoolArray { elems, dims, _id } = cpu_array; + + Self::new(InnerBoolArray::Cpu(elems), dims) + } +} + impl BitwiseArrayBackend for DynFheBoolArrayBackend { fn bitand<'a>( lhs: TensorSlice<'_, Self::Slice<'a>>, diff --git a/tfhe/src/high_level_api/array/gpu/booleans.rs b/tfhe/src/high_level_api/array/gpu/booleans.rs index 2f8b982c53..fdb56275cf 100644 --- a/tfhe/src/high_level_api/array/gpu/booleans.rs +++ b/tfhe/src/high_level_api/array/gpu/booleans.rs @@ -9,7 +9,7 @@ use crate::high_level_api::array::{ArrayBackend, BackendDataContainer, BackendDa use crate::high_level_api::global_state; use crate::integer::gpu::ciphertext::boolean_value::CudaBooleanBlock; use crate::prelude::{FheDecrypt, FheTryEncrypt}; -use crate::{ClientKey, FheBoolId}; +use crate::{ClientKey, FheBool, FheBoolId, Tag}; use rayon::prelude::*; use std::ops::RangeBounds; @@ -72,6 +72,31 @@ impl From> for GpuBooleanOwned { } } +impl From> for GpuFheBoolArray { + fn from(value: Vec) -> Self { + let container = with_cuda_internal_keys(|cuda_keys| { + value + .into_iter() + .map(|val| val.ciphertext.into_gpu(&cuda_keys.streams)) + .collect::>() + }); + + let shape = vec![container.len()]; + Self::new(container, shape) + } +} + +impl From for Vec { + fn from(value: GpuFheBoolArray) -> Self { + value + .into_container() + .0 + .into_iter() + .map(|cuda_bool_block| FheBool::new(cuda_bool_block, Tag::default())) + .collect() + } +} + impl BackendDataContainer for GpuBooleanSlice<'_> { type Backend = GpuFheBoolArrayBackend; diff --git a/tfhe/src/high_level_api/array/gpu/integers.rs b/tfhe/src/high_level_api/array/gpu/integers.rs index 3fc4cbdc8a..50cbdd5825 100644 --- a/tfhe/src/high_level_api/array/gpu/integers.rs +++ b/tfhe/src/high_level_api/array/gpu/integers.rs @@ -24,7 +24,7 @@ use crate::integer::gpu::ciphertext::{ use crate::integer::server_key::radix_parallel::scalar_div_mod::SignedReciprocable; use crate::integer::server_key::{Reciprocable, ScalarMultiplier}; use crate::prelude::{CastInto, FheDecrypt, FheTryEncrypt}; -use crate::{ClientKey, Error}; +use crate::{ClientKey, Error, FheInt, FheUint, Tag}; use rayon::prelude::*; use std::marker::PhantomData; use std::ops::RangeBounds; @@ -60,6 +60,54 @@ where } } +impl From>> for GpuFheUintArray { + fn from(value: Vec>) -> Self { + let container = with_cuda_internal_keys(|cuda_keys| { + value + .into_iter() + .map(|elem| elem.ciphertext.into_gpu(&cuda_keys.streams)) + .collect::>() + }); + let shape = vec![container.len()]; + Self::new(container, shape) + } +} + +impl From> for Vec> { + fn from(value: GpuFheUintArray) -> Self { + value + .into_container() + .0 + .into_iter() + .map(|elem| FheUint::new(elem, Tag::default())) + .collect() + } +} + +impl From>> for GpuFheIntArray { + fn from(value: Vec>) -> Self { + let container = with_cuda_internal_keys(|cuda_keys| { + value + .into_iter() + .map(|elem| elem.ciphertext.into_gpu(&cuda_keys.streams)) + .collect::>() + }); + let shape = vec![container.len()]; + Self::new(container, shape) + } +} + +impl From> for Vec> { + fn from(value: GpuFheIntArray) -> Self { + value + .into_container() + .0 + .into_iter() + .map(|elem| FheInt::new(elem, Tag::default())) + .collect() + } +} + impl ArrayBackend for GpuIntegerArrayBackend where T: CudaIntegerRadixCiphertext, diff --git a/tfhe/src/high_level_api/array/tests/mod.rs b/tfhe/src/high_level_api/array/tests/mod.rs index bc0d6ca562..be9c4273b0 100644 --- a/tfhe/src/high_level_api/array/tests/mod.rs +++ b/tfhe/src/high_level_api/array/tests/mod.rs @@ -2,7 +2,11 @@ mod booleans; mod signed; mod unsigned; -use crate::{generate_keys, set_server_key, ClientKey, ConfigBuilder, FheId}; +use crate::prelude::*; +use crate::{ + generate_keys, set_server_key, ClientKey, ConfigBuilder, CpuFheBoolArray, CpuFheInt32Array, + CpuFheUint32Array, FheBool, FheId, FheInt32, FheUint32, +}; #[cfg(feature = "gpu")] use crate::{Config, CudaServerKey}; use rand::distributions::{Distribution, Standard}; @@ -11,6 +15,8 @@ use std::fmt::Debug; use crate::array::traits::IOwnedArray; use crate::array::ClearArray; +#[cfg(feature = "gpu")] +use crate::array::{GpuFheBoolArray, GpuFheInt32Array, GpuFheUint32Array}; use crate::high_level_api::array::{FheBackendArray, FheBackendArraySlice}; use crate::prelude::{FheDecrypt, FheTryEncrypt}; use std::ops::{BitAnd, BitOr, BitXor}; @@ -293,3 +299,132 @@ where assert_eq!(result, expected_result); } } + +#[cfg(feature = "gpu")] +#[test] +fn test_gpu_conversions() { + let ck = setup_default_gpu(); + + let num_values = 50; + + // Vec <=> GpuFheUint + { + let clears = draw_random_values::(num_values); + let encrypted = clears + .iter() + .map(|v| FheUint32::encrypt(*v, &ck)) + .collect::>(); + + let gpu_array = GpuFheUint32Array::from(encrypted); + let decrypted: Vec = gpu_array.decrypt(&ck); + assert_eq!(decrypted, clears); + + let encrypted = Vec::::from(gpu_array); + let decrypted: Vec = encrypted + .iter() + .map(|fheuint| fheuint.decrypt(&ck)) + .collect(); + assert_eq!(decrypted, clears); + } + + // Vec <=> GpuFheInt + { + let clears = draw_random_values::(num_values); + let encrypted = clears + .iter() + .map(|v| FheInt32::encrypt(*v, &ck)) + .collect::>(); + + let gpu_array = GpuFheInt32Array::from(encrypted); + let decrypted: Vec = gpu_array.decrypt(&ck); + assert_eq!(decrypted, clears); + + let encrypted = Vec::::from(gpu_array); + let decrypted: Vec = encrypted.iter().map(|fheint| fheint.decrypt(&ck)).collect(); + assert_eq!(decrypted, clears); + } + + // Vec <=> GpuFheBool + { + let clears = draw_random_values::(num_values); + let encrypted = clears + .iter() + .map(|v| FheBool::encrypt(*v, &ck)) + .collect::>(); + + let gpu_array = GpuFheBoolArray::from(encrypted); + let decrypted: Vec = gpu_array.decrypt(&ck); + assert_eq!(decrypted, clears); + + let encrypted = Vec::::from(gpu_array); + let decrypted: Vec = encrypted + .iter() + .map(|fhebool| fhebool.decrypt(&ck)) + .collect(); + assert_eq!(decrypted, clears); + } +} + +#[test] +fn test_cpu_conversions() { + let ck = setup_default_cpu(); + + let num_values = 50; + + // Vec <=> CpuFheUint + { + let clears = draw_random_values::(num_values); + let encrypted = clears + .iter() + .map(|v| FheUint32::encrypt(*v, &ck)) + .collect::>(); + + let cpu_array = CpuFheUint32Array::from(encrypted); + let decrypted: Vec = cpu_array.decrypt(&ck); + assert_eq!(decrypted, clears); + + let encrypted = Vec::::from(cpu_array); + let decrypted: Vec = encrypted + .iter() + .map(|fheuint| fheuint.decrypt(&ck)) + .collect(); + assert_eq!(decrypted, clears); + } + + // Vec <=> CpuFheInt + { + let clears = draw_random_values::(num_values); + let encrypted = clears + .iter() + .map(|v| FheInt32::encrypt(*v, &ck)) + .collect::>(); + + let cpu_array = CpuFheInt32Array::from(encrypted); + let decrypted: Vec = cpu_array.decrypt(&ck); + assert_eq!(decrypted, clears); + + let encrypted = Vec::::from(cpu_array); + let decrypted: Vec = encrypted.iter().map(|fheint| fheint.decrypt(&ck)).collect(); + assert_eq!(decrypted, clears); + } + + // Vec <=> CpuFheBool + { + let clears = draw_random_values::(num_values); + let encrypted = clears + .iter() + .map(|v| FheBool::encrypt(*v, &ck)) + .collect::>(); + + let cpu_array = CpuFheBoolArray::from(encrypted); + let decrypted: Vec = cpu_array.decrypt(&ck); + assert_eq!(decrypted, clears); + + let encrypted = Vec::::from(cpu_array); + let decrypted: Vec = encrypted + .iter() + .map(|fhebool| fhebool.decrypt(&ck)) + .collect(); + assert_eq!(decrypted, clears); + } +} diff --git a/tfhe/src/high_level_api/booleans/inner.rs b/tfhe/src/high_level_api/booleans/inner.rs index 981acffe2f..f810e463e7 100644 --- a/tfhe/src/high_level_api/booleans/inner.rs +++ b/tfhe/src/high_level_api/booleans/inner.rs @@ -25,7 +25,7 @@ use crate::integer::hpu::ciphertext::HpuRadixCiphertext; pub(in crate::high_level_api) enum InnerBoolean { Cpu(BooleanBlock), #[cfg(feature = "gpu")] - Cuda(crate::integer::gpu::ciphertext::boolean_value::CudaBooleanBlock), + Cuda(CudaBooleanBlock), #[cfg(feature = "hpu")] Hpu(HpuRadixCiphertext), } @@ -220,7 +220,6 @@ impl InnerBoolean { &mut cuda_ct.0 } - #[cfg(feature = "gpu")] pub(crate) fn into_cpu(self) -> BooleanBlock { match self { Self::Cpu(cpu_ct) => cpu_ct,