tfhe/high_level_api/array/dynamic/
booleans.rs

1//! This module contains the implementation of the FheBool array backend
2//! where the location of the values and computations can be changed/selected at runtime
3use super::super::cpu::{CpuFheBoolArrayBackend, FheBoolId};
4use super::super::helpers::{create_sub_mut_slice_with_bound, create_sub_slice_with_bound};
5use super::super::traits::{
6    ArrayBackend, BackendDataContainer, BackendDataContainerMut, BitwiseArrayBackend,
7    ClearBitwiseArrayBackend,
8};
9use super::super::{FheBackendArray, FheBackendArraySlice, FheBackendArraySliceMut};
10
11use crate::array::traits::TensorSlice;
12use crate::integer::BooleanBlock;
13use crate::prelude::{FheDecrypt, FheTryEncrypt};
14use crate::{ClientKey, Device};
15use std::borrow::{Borrow, Cow};
16use std::ops::RangeBounds;
17
18pub type FheBoolArray = FheBackendArray<DynFheBoolArrayBackend, FheBoolId>;
19pub type FheBoolSlice<'a> = FheBackendArraySlice<'a, DynFheBoolArrayBackend, FheBoolId>;
20pub type FheBoolSliceMut<'a> = FheBackendArraySliceMut<'a, DynFheBoolArrayBackend, FheBoolId>;
21
22pub struct DynFheBoolArrayBackend;
23
24impl ArrayBackend for DynFheBoolArrayBackend {
25    type Slice<'a>
26        = InnerBoolSlice<'a>
27    where
28        Self: 'a;
29    type SliceMut<'a>
30        = InnerBoolSliceMut<'a>
31    where
32        Self: 'a;
33    type Owned = InnerBoolArray;
34}
35
36impl BitwiseArrayBackend for DynFheBoolArrayBackend {
37    fn bitand<'a>(
38        lhs: TensorSlice<'_, Self::Slice<'a>>,
39        rhs: TensorSlice<'_, Self::Slice<'a>>,
40    ) -> Self::Owned {
41        dispatch_binary_op(&lhs, &rhs, CpuFheBoolArrayBackend::bitand)
42    }
43
44    fn bitor<'a>(
45        lhs: TensorSlice<'_, Self::Slice<'a>>,
46        rhs: TensorSlice<'_, Self::Slice<'a>>,
47    ) -> Self::Owned {
48        dispatch_binary_op(&lhs, &rhs, CpuFheBoolArrayBackend::bitor)
49    }
50
51    fn bitxor<'a>(
52        lhs: TensorSlice<'_, Self::Slice<'a>>,
53        rhs: TensorSlice<'_, Self::Slice<'a>>,
54    ) -> Self::Owned {
55        dispatch_binary_op(&lhs, &rhs, CpuFheBoolArrayBackend::bitxor)
56    }
57
58    fn bitnot(lhs: TensorSlice<'_, Self::Slice<'_>>) -> Self::Owned {
59        dispatch_unary_op(&lhs, CpuFheBoolArrayBackend::bitnot)
60    }
61}
62
63impl ClearBitwiseArrayBackend<bool> for DynFheBoolArrayBackend {
64    fn bitand_slice(
65        lhs: TensorSlice<'_, Self::Slice<'_>>,
66        rhs: TensorSlice<'_, &'_ [bool]>,
67    ) -> Self::Owned {
68        dispatch_binary_scalar_op(&lhs, &rhs, CpuFheBoolArrayBackend::bitand_slice)
69    }
70
71    fn bitor_slice(
72        lhs: TensorSlice<'_, Self::Slice<'_>>,
73        rhs: TensorSlice<'_, &'_ [bool]>,
74    ) -> Self::Owned {
75        dispatch_binary_scalar_op(&lhs, &rhs, CpuFheBoolArrayBackend::bitor_slice)
76    }
77
78    fn bitxor_slice(
79        lhs: TensorSlice<'_, Self::Slice<'_>>,
80        rhs: TensorSlice<'_, &'_ [bool]>,
81    ) -> Self::Owned {
82        dispatch_binary_scalar_op(&lhs, &rhs, CpuFheBoolArrayBackend::bitxor_slice)
83    }
84}
85
86#[inline]
87fn dispatch_binary_op<CpuFn>(
88    lhs: &TensorSlice<'_, InnerBoolSlice<'_>>,
89    rhs: &TensorSlice<'_, InnerBoolSlice<'_>>,
90    cpu_fn: CpuFn,
91) -> InnerBoolArray
92where
93    CpuFn: for<'a> Fn(
94        TensorSlice<'_, &'a [BooleanBlock]>,
95        TensorSlice<'_, &'a [BooleanBlock]>,
96    ) -> Vec<BooleanBlock>,
97{
98    match crate::high_level_api::global_state::device_of_internal_keys() {
99        Some(Device::Cpu) => {
100            let lhs_cpu_cow = lhs.slice.on_cpu();
101            let rhs_cpu_cow = rhs.slice.on_cpu();
102
103            let lhs_cpu_slice: &[BooleanBlock] = lhs_cpu_cow.borrow();
104            let rhs_cpu_slice: &[BooleanBlock] = rhs_cpu_cow.borrow();
105
106            let result = cpu_fn(
107                TensorSlice::new(lhs_cpu_slice, lhs.dims),
108                TensorSlice::new(rhs_cpu_slice, rhs.dims),
109            );
110            InnerBoolArray::Cpu(result)
111        }
112        #[cfg(feature = "gpu")]
113        Some(Device::CudaGpu) => {
114            panic!("Not supported by Cuda devices")
115        }
116        #[cfg(feature = "hpu")]
117        Some(Device::Hpu) => {
118            panic!("Not supported by Hpu devices")
119        }
120        None => {
121            panic!("{}", crate::high_level_api::errors::UninitializedServerKey);
122        }
123    }
124}
125
126#[inline]
127fn dispatch_unary_op<CpuFn>(
128    lhs: &TensorSlice<'_, InnerBoolSlice<'_>>,
129    cpu_fn: CpuFn,
130) -> InnerBoolArray
131where
132    CpuFn: for<'a> Fn(TensorSlice<'_, &'a [BooleanBlock]>) -> Vec<BooleanBlock>,
133{
134    match crate::high_level_api::global_state::device_of_internal_keys() {
135        Some(Device::Cpu) => {
136            let lhs_cpu_cow = lhs.slice.on_cpu();
137
138            let lhs_cpu_slice: &[BooleanBlock] = lhs_cpu_cow.borrow();
139
140            let result = cpu_fn(TensorSlice::new(lhs_cpu_slice, lhs.dims));
141            InnerBoolArray::Cpu(result)
142        }
143        #[cfg(feature = "gpu")]
144        Some(Device::CudaGpu) => {
145            panic!("Not supported by Cuda devices")
146        }
147        #[cfg(feature = "hpu")]
148        Some(Device::Hpu) => {
149            panic!("Not supported by Hpu devices")
150        }
151        None => {
152            panic!("{}", crate::high_level_api::errors::UninitializedServerKey);
153        }
154    }
155}
156
157#[inline]
158fn dispatch_binary_scalar_op<CpuFn>(
159    lhs: &TensorSlice<'_, InnerBoolSlice<'_>>,
160    rhs: &TensorSlice<'_, &'_ [bool]>,
161    cpu_fn: CpuFn,
162) -> InnerBoolArray
163where
164    CpuFn: for<'a> Fn(
165        TensorSlice<'_, &'a [BooleanBlock]>,
166        TensorSlice<'_, &'a [bool]>,
167    ) -> Vec<BooleanBlock>,
168{
169    match crate::high_level_api::global_state::device_of_internal_keys() {
170        Some(Device::Cpu) => {
171            let lhs_cpu_cow = lhs.slice.on_cpu();
172
173            let lhs_cpu_slice: &[BooleanBlock] = lhs_cpu_cow.borrow();
174
175            let result = cpu_fn(
176                TensorSlice::new(lhs_cpu_slice, lhs.dims),
177                TensorSlice::new(rhs.slice, rhs.dims),
178            );
179            InnerBoolArray::Cpu(result)
180        }
181        #[cfg(feature = "gpu")]
182        Some(Device::CudaGpu) => {
183            panic!("Not supported by Cuda devices")
184        }
185        #[cfg(feature = "hpu")]
186        Some(Device::Hpu) => {
187            panic!("Not supported by Hpu devices")
188        }
189        None => {
190            panic!("{}", crate::high_level_api::errors::UninitializedServerKey);
191        }
192    }
193}
194
195#[derive(Clone)]
196pub enum InnerBoolArray {
197    Cpu(Vec<BooleanBlock>),
198}
199
200impl BackendDataContainer for InnerBoolArray {
201    type Backend = DynFheBoolArrayBackend;
202
203    fn len(&self) -> usize {
204        match self {
205            Self::Cpu(cpu_array) => cpu_array.len(),
206        }
207    }
208
209    fn as_sub_slice(
210        &self,
211        range: impl RangeBounds<usize>,
212    ) -> <Self::Backend as ArrayBackend>::Slice<'_> {
213        match self {
214            Self::Cpu(cpu_vec) => {
215                InnerBoolSlice::Cpu(create_sub_slice_with_bound(cpu_vec.as_slice(), range))
216            }
217        }
218    }
219
220    fn into_owned(self) -> <Self::Backend as ArrayBackend>::Owned {
221        self
222    }
223}
224
225impl BackendDataContainerMut for InnerBoolArray {
226    fn as_sub_slice_mut(
227        &mut self,
228        range: impl RangeBounds<usize>,
229    ) -> <Self::Backend as ArrayBackend>::SliceMut<'_> {
230        match self {
231            Self::Cpu(cpu_vec) => InnerBoolSliceMut::Cpu(create_sub_mut_slice_with_bound(
232                cpu_vec.as_mut_slice(),
233                range,
234            )),
235        }
236    }
237}
238
239#[derive(Copy, Clone)]
240pub enum InnerBoolSlice<'a> {
241    Cpu(&'a [BooleanBlock]),
242}
243
244impl InnerBoolSlice<'_> {
245    fn on_cpu(&self) -> Cow<'_, [BooleanBlock]> {
246        match self {
247            InnerBoolSlice::Cpu(cpu_slice) => Cow::Borrowed(cpu_slice),
248        }
249    }
250}
251
252impl BackendDataContainer for InnerBoolSlice<'_> {
253    type Backend = DynFheBoolArrayBackend;
254
255    fn len(&self) -> usize {
256        match self {
257            InnerBoolSlice::Cpu(cpu_slice) => cpu_slice.len(),
258        }
259    }
260
261    fn as_sub_slice(
262        &self,
263        range: impl RangeBounds<usize>,
264    ) -> <Self::Backend as ArrayBackend>::Slice<'_> {
265        match self {
266            InnerBoolSlice::Cpu(cpu_slice) => {
267                InnerBoolSlice::Cpu(create_sub_slice_with_bound(*cpu_slice, range))
268            }
269        }
270    }
271
272    fn into_owned(self) -> <Self::Backend as ArrayBackend>::Owned {
273        match self {
274            InnerBoolSlice::Cpu(cpu_slice) => InnerBoolArray::Cpu(cpu_slice.to_vec()),
275        }
276    }
277}
278
279pub enum InnerBoolSliceMut<'a> {
280    Cpu(&'a mut [BooleanBlock]),
281}
282
283impl BackendDataContainer for InnerBoolSliceMut<'_> {
284    type Backend = DynFheBoolArrayBackend;
285
286    fn len(&self) -> usize {
287        match self {
288            InnerBoolSliceMut::Cpu(cpu_slice) => cpu_slice.len(),
289        }
290    }
291
292    fn as_sub_slice(
293        &self,
294        range: impl RangeBounds<usize>,
295    ) -> <Self::Backend as ArrayBackend>::Slice<'_> {
296        match self {
297            Self::Cpu(cpu_slice) => {
298                InnerBoolSlice::Cpu(create_sub_slice_with_bound(*cpu_slice, range))
299            }
300        }
301    }
302
303    fn into_owned(self) -> <Self::Backend as ArrayBackend>::Owned {
304        match self {
305            Self::Cpu(cpu_slice) => InnerBoolArray::Cpu(cpu_slice.to_vec()),
306        }
307    }
308}
309
310impl BackendDataContainerMut for InnerBoolSliceMut<'_> {
311    fn as_sub_slice_mut(
312        &mut self,
313        range: impl RangeBounds<usize>,
314    ) -> <Self::Backend as ArrayBackend>::SliceMut<'_> {
315        match self {
316            InnerBoolSliceMut::Cpu(cpu_slice) => {
317                InnerBoolSliceMut::Cpu(create_sub_mut_slice_with_bound(*cpu_slice, range))
318            }
319        }
320    }
321}
322
323impl<'a> FheTryEncrypt<&'a [bool], ClientKey> for FheBoolArray {
324    type Error = crate::Error;
325
326    fn try_encrypt(value: &'a [bool], key: &ClientKey) -> Result<Self, Self::Error> {
327        let cpu_array = crate::CpuFheBoolArray::try_encrypt(value, key)?;
328        let inner = InnerBoolArray::Cpu(cpu_array.into_container());
329        // TODO move to default device
330        Ok(Self::new(inner, vec![value.len()]))
331    }
332}
333
334impl FheDecrypt<Vec<bool>> for FheBoolArray {
335    fn decrypt(&self, key: &ClientKey) -> Vec<bool> {
336        let slice = self.elems.as_slice();
337        let cpu_cow = slice.on_cpu();
338        let cpu_slice = cpu_cow.as_ref();
339
340        crate::CpuFheBoolSlice::<'_>::new(cpu_slice, self.dims.clone()).decrypt(key)
341    }
342}
tfhe/high_level_api/array/dynamic/booleans.rs

tfhe/high_level_api/array/dynamic/
booleans.rs
