feat(gpu): add gpu array type in hl api

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

@@ -0,0 +1,256 @@
+//! This module contains the implementation of the FheBool array backend
+//! where the values and computations are always done on CPU
+use super::super::helpers::{create_sub_mut_slice_with_bound, create_sub_slice_with_bound};
+use super::super::traits::{BitwiseArrayBackend, ClearBitwiseArrayBackend};
+use crate::array::traits::TensorSlice;
+use crate::high_level_api::array::{ArrayBackend, BackendDataContainer, BackendDataContainerMut};
+use crate::high_level_api::global_state;
+use crate::high_level_api::global_state::with_thread_local_cuda_streams;
+use crate::integer::gpu::ciphertext::boolean_value::CudaBooleanBlock;
+use crate::prelude::{FheDecrypt, FheTryEncrypt};
+use crate::{ClientKey, FheBoolId};
+use rayon::prelude::*;
+use std::ops::RangeBounds;
+
+pub struct GpuFheBoolArrayBackend;
+
+pub type GpuFheBoolArray = super::super::FheBackendArray<GpuFheBoolArrayBackend, FheBoolId>;
+pub type GpuFheBoolSlice<'a> =
+    super::super::FheBackendArraySlice<'a, GpuFheBoolArrayBackend, FheBoolId>;
+pub type GpuFheBoolSliceMut<'a> =
+    super::super::FheBackendArraySliceMut<'a, GpuFheBoolArrayBackend, FheBoolId>;
+
+impl ArrayBackend for GpuFheBoolArrayBackend {
+    type Slice<'a>
+        = &'a [CudaBooleanBlock]
+    where
+        Self: 'a;
+    type SliceMut<'a>
+        = &'a mut [CudaBooleanBlock]
+    where
+        Self: 'a;
+    type Owned = Vec<CudaBooleanBlock>;
+}
+
+impl<'a> BackendDataContainer for &'a [CudaBooleanBlock] {
+    type Backend = GpuFheBoolArrayBackend;
+
+    fn len(&self) -> usize {
+        <[CudaBooleanBlock]>::len(self)
+    }
+
+    fn as_sub_slice(
+        &self,
+        range: impl RangeBounds<usize>,
+    ) -> <Self::Backend as ArrayBackend>::Slice<'_> {
+        create_sub_slice_with_bound(*self, range)
+    }
+
+    fn into_owned(self) -> <Self::Backend as ArrayBackend>::Owned {
+        self.to_vec()
+    }
+}
+
+impl<'a> BackendDataContainer for &'a mut [CudaBooleanBlock] {
+    type Backend = GpuFheBoolArrayBackend;
+
+    fn len(&self) -> usize {
+        <[CudaBooleanBlock]>::len(self)
+    }
+
+    fn as_sub_slice(
+        &self,
+        range: impl RangeBounds<usize>,
+    ) -> <Self::Backend as ArrayBackend>::Slice<'_> {
+        create_sub_slice_with_bound(*self, range)
+    }
+
+    fn into_owned(self) -> <Self::Backend as ArrayBackend>::Owned {
+        self.to_vec()
+    }
+}
+
+impl<'a> BackendDataContainerMut for &'a mut [CudaBooleanBlock] {
+    fn as_sub_slice_mut(
+        &mut self,
+        range: impl RangeBounds<usize>,
+    ) -> <Self::Backend as ArrayBackend>::SliceMut<'_> {
+        create_sub_mut_slice_with_bound(*self, range)
+    }
+}
+
+impl BackendDataContainer for Vec<CudaBooleanBlock> {
+    type Backend = GpuFheBoolArrayBackend;
+
+    fn len(&self) -> usize {
+        self.len()
+    }
+
+    fn as_sub_slice(
+        &self,
+        range: impl RangeBounds<usize>,
+    ) -> <Self::Backend as ArrayBackend>::Slice<'_> {
+        create_sub_slice_with_bound(self, range)
+    }
+
+    fn into_owned(self) -> <Self::Backend as ArrayBackend>::Owned {
+        self
+    }
+}
+
+impl BackendDataContainerMut for Vec<CudaBooleanBlock> {
+    fn as_sub_slice_mut(
+        &mut self,
+        range: impl RangeBounds<usize>,
+    ) -> <Self::Backend as ArrayBackend>::SliceMut<'_> {
+        create_sub_mut_slice_with_bound(self.as_mut_slice(), range)
+    }
+}
+
+impl BitwiseArrayBackend for GpuFheBoolArrayBackend {
+    fn bitand<'a>(
+        lhs: TensorSlice<'_, Self::Slice<'a>>,
+        rhs: TensorSlice<'_, Self::Slice<'a>>,
+    ) -> Self::Owned {
+        global_state::with_cuda_internal_keys(|cuda_key| {
+            with_thread_local_cuda_streams(|streams| {
+                lhs.par_iter()
+                    .zip(rhs.par_iter())
+                    .map(|(lhs, rhs)| CudaBooleanBlock(cuda_key.bitand(&lhs.0, &rhs.0, streams)))
+                    .collect()
+            })
+        })
+    }
+
+    fn bitor<'a>(
+        lhs: TensorSlice<'_, Self::Slice<'a>>,
+        rhs: TensorSlice<'_, Self::Slice<'a>>,
+    ) -> Self::Owned {
+        global_state::with_cuda_internal_keys(|cuda_key| {
+            with_thread_local_cuda_streams(|streams| {
+                lhs.par_iter()
+                    .zip(rhs.par_iter())
+                    .map(|(lhs, rhs)| CudaBooleanBlock(cuda_key.bitor(&lhs.0, &rhs.0, streams)))
+                    .collect()
+            })
+        })
+    }
+
+    fn bitxor<'a>(
+        lhs: TensorSlice<'_, Self::Slice<'a>>,
+        rhs: TensorSlice<'_, Self::Slice<'a>>,
+    ) -> Self::Owned {
+        global_state::with_cuda_internal_keys(|cuda_key| {
+            with_thread_local_cuda_streams(|streams| {
+                lhs.par_iter()
+                    .zip(rhs.par_iter())
+                    .map(|(lhs, rhs)| CudaBooleanBlock(cuda_key.bitxor(&lhs.0, &rhs.0, streams)))
+                    .collect()
+            })
+        })
+    }
+
+    fn bitnot(lhs: TensorSlice<'_, Self::Slice<'_>>) -> Self::Owned {
+        global_state::with_cuda_internal_keys(|cuda_key| {
+            with_thread_local_cuda_streams(|streams| {
+                lhs.par_iter()
+                    .map(|lhs| CudaBooleanBlock(cuda_key.bitnot(&lhs.0, streams)))
+                    .collect()
+            })
+        })
+    }
+}
+
+impl ClearBitwiseArrayBackend<bool> for GpuFheBoolArrayBackend {
+    fn bitand_slice(
+        lhs: TensorSlice<'_, Self::Slice<'_>>,
+        rhs: TensorSlice<'_, &'_ [bool]>,
+    ) -> Self::Owned {
+        global_state::with_cuda_internal_keys(|cuda_key| {
+            with_thread_local_cuda_streams(|streams| {
+                lhs.par_iter()
+                    .zip(rhs.par_iter().copied())
+                    .map(|(lhs, rhs)| {
+                        CudaBooleanBlock(cuda_key.scalar_bitand(&lhs.0, rhs as u8, streams))
+                    })
+                    .collect()
+            })
+        })
+    }
+
+    fn bitor_slice(
+        lhs: TensorSlice<'_, Self::Slice<'_>>,
+        rhs: TensorSlice<'_, &'_ [bool]>,
+    ) -> Self::Owned {
+        global_state::with_cuda_internal_keys(|cuda_key| {
+            with_thread_local_cuda_streams(|streams| {
+                lhs.par_iter()
+                    .zip(rhs.par_iter().copied())
+                    .map(|(lhs, rhs)| {
+                        CudaBooleanBlock(cuda_key.scalar_bitor(&lhs.0, rhs as u8, streams))
+                    })
+                    .collect()
+            })
+        })
+    }
+
+    fn bitxor_slice(
+        lhs: TensorSlice<'_, Self::Slice<'_>>,
+        rhs: TensorSlice<'_, &'_ [bool]>,
+    ) -> Self::Owned {
+        global_state::with_cuda_internal_keys(|cuda_key| {
+            with_thread_local_cuda_streams(|streams| {
+                lhs.par_iter()
+                    .zip(rhs.par_iter().copied())
+                    .map(|(lhs, rhs)| {
+                        CudaBooleanBlock(cuda_key.scalar_bitxor(&lhs.0, rhs as u8, streams))
+                    })
+                    .collect()
+            })
+        })
+    }
+}
+
+impl FheTryEncrypt<&[bool], ClientKey> for GpuFheBoolArray {
+    type Error = crate::Error;
+
+    fn try_encrypt(values: &[bool], cks: &ClientKey) -> Result<Self, Self::Error> {
+        let encrypted = with_thread_local_cuda_streams(|streams| {
+            values
+                .iter()
+                .copied()
+                .map(|value| {
+                    CudaBooleanBlock::from_boolean_block(&cks.key.key.encrypt_bool(value), streams)
+                })
+                .collect::<Vec<_>>()
+        });
+        Ok(Self::new(encrypted, vec![values.len()]))
+    }
+}
+
+impl<'a> FheDecrypt<Vec<bool>> for GpuFheBoolSlice<'a> {
+    fn decrypt(&self, key: &ClientKey) -> Vec<bool> {
+        with_thread_local_cuda_streams(|streams| {
+            self.elems
+                .iter()
+                .map(|encrypted_value| {
+                    key.key
+                        .key
+                        .decrypt_bool(&encrypted_value.to_boolean_block(streams))
+                })
+                .collect()
+        })
+    }
+}
+
+impl<'a> FheDecrypt<Vec<bool>> for GpuFheBoolSliceMut<'a> {
+    fn decrypt(&self, key: &ClientKey) -> Vec<bool> {
+        self.as_slice().decrypt(key)
+    }
+}
+
+impl FheDecrypt<Vec<bool>> for GpuFheBoolArray {
+    fn decrypt(&self, key: &ClientKey) -> Vec<bool> {
+        self.as_slice().decrypt(key)
+    }
+}