Skip to content

Commit 3fdc289

Browse files
tmontaigutmontaigu
committed
chore(hlapi): Add array conversion from/to Vec<FheType>
Add `From` impl to allow conversion from Vec<FheType> like Vec<FheUint32> to Cpu/Gpu array.
1 parent a5c876f commit 3fdc289

File tree

7 files changed

+316
-8
lines changed

7 files changed

+316
-8
lines changed

‎tfhe/src/high_level_api/array/cpu/booleans.rs‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use crate::high_level_api::array::{ArrayBackend, BackendDataContainer, BackendDa
77
use crate::high_level_api::global_state;
88
use crate::integer::BooleanBlock;
99
use crate::prelude::{FheDecrypt, FheTryEncrypt};
10-
use crate::{ClientKey, FheId};
10+
use crate::{ClientKey, FheBool, FheId, Tag};
1111
use rayon::prelude::*;
1212
use std::ops::RangeBounds;
1313

@@ -36,6 +36,28 @@ impl ArrayBackend for CpuFheBoolArrayBackend {
3636
type Owned = Vec<BooleanBlock>;
3737
}
3838

39+
impl From<Vec<FheBool>> for CpuFheBoolArray {
40+
fn from(value: Vec<FheBool>) -> Self {
41+
let vec = value
42+
.into_iter()
43+
.map(|b| BooleanBlock::new_unchecked(b.into_raw_parts()))
44+
.collect::<Vec<_>>();
45+
46+
let shape = vec![vec.len()];
47+
Self::new(vec, shape)
48+
}
49+
}
50+
51+
impl From<CpuFheBoolArray> for Vec<FheBool> {
52+
fn from(value: CpuFheBoolArray) -> Self {
53+
value
54+
.into_container()
55+
.into_iter()
56+
.map(|boolean_block| FheBool::new(boolean_block, Tag::default()))
57+
.collect()
58+
}
59+
}
60+
3961
impl BackendDataContainer for &[BooleanBlock] {
4062
type Backend = CpuFheBoolArrayBackend;
4163

‎tfhe/src/high_level_api/array/cpu/integers.rs‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use crate::integer::server_key::radix_parallel::scalar_div_mod::SignedReciprocab
1919
use crate::integer::server_key::{Reciprocable, ScalarMultiplier};
2020
use crate::integer::{IntegerRadixCiphertext, RadixCiphertext, SignedRadixCiphertext};
2121
use crate::prelude::{FheDecrypt, FheTryEncrypt};
22-
use crate::{ClientKey, Error};
22+
use crate::{ClientKey, Error, FheInt, FheUint, Tag};
2323
use rayon::prelude::*;
2424
use std::marker::PhantomData;
2525
use std::ops::RangeBounds;
@@ -54,6 +54,48 @@ where
5454
type Owned = Vec<T>;
5555
}
5656

57+
impl<Id: FheUintId> From<Vec<FheUint<Id>>> for CpuFheUintArray<Id> {
58+
fn from(value: Vec<FheUint<Id>>) -> Self {
59+
let vec: Vec<_> = value
60+
.into_iter()
61+
.map(|uint| uint.into_raw_parts().0)
62+
.collect();
63+
let shape = vec![vec.len()];
64+
Self::new(vec, shape)
65+
}
66+
}
67+
68+
impl<Id: FheUintId> From<CpuFheUintArray<Id>> for Vec<FheUint<Id>> {
69+
fn from(value: CpuFheUintArray<Id>) -> Self {
70+
value
71+
.into_container()
72+
.into_iter()
73+
.map(|radix| FheUint::new(radix, Tag::default()))
74+
.collect()
75+
}
76+
}
77+
78+
impl<Id: FheIntId> From<Vec<FheInt<Id>>> for CpuFheIntArray<Id> {
79+
fn from(value: Vec<FheInt<Id>>) -> Self {
80+
let vec: Vec<_> = value
81+
.into_iter()
82+
.map(|uint| uint.into_raw_parts().0)
83+
.collect();
84+
let shape = vec![vec.len()];
85+
Self::new(vec, shape)
86+
}
87+
}
88+
89+
impl<Id: FheIntId> From<CpuFheIntArray<Id>> for Vec<FheInt<Id>> {
90+
fn from(value: CpuFheIntArray<Id>) -> Self {
91+
value
92+
.into_container()
93+
.into_iter()
94+
.map(|radix| FheInt::new(radix, Tag::default()))
95+
.collect()
96+
}
97+
}
98+
5799
#[inline]
58100
#[track_caller]
59101
fn par_map_sks_op_on_pair_of_elements<'a, T, F>(

‎tfhe/src/high_level_api/array/dynamic/booleans.rs‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use super::super::{FheBackendArray, FheBackendArraySlice, FheBackendArraySliceMu
1111
use crate::array::traits::TensorSlice;
1212
use crate::integer::BooleanBlock;
1313
use crate::prelude::{FheDecrypt, FheTryEncrypt};
14-
use crate::{ClientKey, Device};
14+
use crate::{ClientKey, CpuFheBoolArray, Device, FheBool};
1515
use std::borrow::{Borrow, Cow};
1616
use std::ops::RangeBounds;
1717

@@ -33,6 +33,43 @@ impl ArrayBackend for DynFheBoolArrayBackend {
3333
type Owned = InnerBoolArray;
3434
}
3535

36+
impl TryFrom<Vec<FheBool>> for FheBoolArray {
37+
type Error = crate::Error;
38+
39+
fn try_from(values: Vec<FheBool>) -> Result<Self, Self::Error> {
40+
if values.is_empty() {
41+
return Ok(Self::new(InnerBoolArray::Cpu(vec![]), vec![0]));
42+
}
43+
44+
let shape = vec![values.len()];
45+
let device_of_first = values[0].current_device();
46+
let inner = match device_of_first {
47+
Device::Cpu => {
48+
let new_values = values
49+
.into_iter()
50+
.map(|value| value.ciphertext.into_cpu())
51+
.collect::<Vec<_>>();
52+
53+
InnerBoolArray::Cpu(new_values)
54+
}
55+
#[cfg(feature = "gpu")]
56+
Device::CudaGpu => return crate::error!("Array do not support GPU"),
57+
#[cfg(feature = "hpu")]
58+
Device::Hpu => return crate::error!("Array do not support HPU"),
59+
};
60+
61+
Ok(Self::new(inner, shape))
62+
}
63+
}
64+
65+
impl From<CpuFheBoolArray> for FheBoolArray {
66+
fn from(cpu_array: CpuFheBoolArray) -> Self {
67+
let CpuFheBoolArray { elems, dims, _id } = cpu_array;
68+
69+
Self::new(InnerBoolArray::Cpu(elems), dims)
70+
}
71+
}
72+
3673
impl BitwiseArrayBackend for DynFheBoolArrayBackend {
3774
fn bitand<'a>(
3875
lhs: TensorSlice<'_, Self::Slice<'a>>,

‎tfhe/src/high_level_api/array/gpu/booleans.rs‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use crate::high_level_api::array::{ArrayBackend, BackendDataContainer, BackendDa
99
use crate::high_level_api::global_state;
1010
use crate::integer::gpu::ciphertext::boolean_value::CudaBooleanBlock;
1111
use crate::prelude::{FheDecrypt, FheTryEncrypt};
12-
use crate::{ClientKey, FheBoolId};
12+
use crate::{ClientKey, FheBool, FheBoolId, Tag};
1313
use rayon::prelude::*;
1414
use std::ops::RangeBounds;
1515

@@ -72,6 +72,31 @@ impl From<Vec<CudaBooleanBlock>> for GpuBooleanOwned {
7272
}
7373
}
7474

75+
impl From<Vec<FheBool>> for GpuFheBoolArray {
76+
fn from(value: Vec<FheBool>) -> Self {
77+
let container = with_cuda_internal_keys(|cuda_keys| {
78+
value
79+
.into_iter()
80+
.map(|val| val.ciphertext.into_gpu(&cuda_keys.streams))
81+
.collect::<Vec<_>>()
82+
});
83+
84+
let shape = vec![container.len()];
85+
Self::new(container, shape)
86+
}
87+
}
88+
89+
impl From<GpuFheBoolArray> for Vec<FheBool> {
90+
fn from(value: GpuFheBoolArray) -> Self {
91+
value
92+
.into_container()
93+
.0
94+
.into_iter()
95+
.map(|cuda_bool_block| FheBool::new(cuda_bool_block, Tag::default()))
96+
.collect()
97+
}
98+
}
99+
75100
impl BackendDataContainer for GpuBooleanSlice<'_> {
76101
type Backend = GpuFheBoolArrayBackend;
77102

‎tfhe/src/high_level_api/array/gpu/integers.rs‎

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ use crate::integer::gpu::ciphertext::{
2424
use crate::integer::server_key::radix_parallel::scalar_div_mod::SignedReciprocable;
2525
use crate::integer::server_key::{Reciprocable, ScalarMultiplier};
2626
use crate::prelude::{CastInto, FheDecrypt, FheTryEncrypt};
27-
use crate::{ClientKey, Error};
27+
use crate::{ClientKey, Error, FheInt, FheUint, Tag};
2828
use rayon::prelude::*;
2929
use std::marker::PhantomData;
3030
use std::ops::RangeBounds;
@@ -60,6 +60,54 @@ where
6060
}
6161
}
6262

63+
impl<Id: FheUintId> From<Vec<FheUint<Id>>> for GpuFheUintArray<Id> {
64+
fn from(value: Vec<FheUint<Id>>) -> Self {
65+
let container = with_cuda_internal_keys(|cuda_keys| {
66+
value
67+
.into_iter()
68+
.map(|elem| elem.ciphertext.into_gpu(&cuda_keys.streams))
69+
.collect::<Vec<_>>()
70+
});
71+
let shape = vec![container.len()];
72+
Self::new(container, shape)
73+
}
74+
}
75+
76+
impl<Id: FheUintId> From<GpuFheUintArray<Id>> for Vec<FheUint<Id>> {
77+
fn from(value: GpuFheUintArray<Id>) -> Self {
78+
value
79+
.into_container()
80+
.0
81+
.into_iter()
82+
.map(|elem| FheUint::new(elem, Tag::default()))
83+
.collect()
84+
}
85+
}
86+
87+
impl<Id: FheIntId> From<Vec<FheInt<Id>>> for GpuFheIntArray<Id> {
88+
fn from(value: Vec<FheInt<Id>>) -> Self {
89+
let container = with_cuda_internal_keys(|cuda_keys| {
90+
value
91+
.into_iter()
92+
.map(|elem| elem.ciphertext.into_gpu(&cuda_keys.streams))
93+
.collect::<Vec<_>>()
94+
});
95+
let shape = vec![container.len()];
96+
Self::new(container, shape)
97+
}
98+
}
99+
100+
impl<Id: FheIntId> From<GpuFheIntArray<Id>> for Vec<FheInt<Id>> {
101+
fn from(value: GpuFheIntArray<Id>) -> Self {
102+
value
103+
.into_container()
104+
.0
105+
.into_iter()
106+
.map(|elem| FheInt::new(elem, Tag::default()))
107+
.collect()
108+
}
109+
}
110+
63111
impl<T> ArrayBackend for GpuIntegerArrayBackend<T>
64112
where
65113
T: CudaIntegerRadixCiphertext,

0 commit comments

Comments
 (0)