chore(gpu): add vectorized function for bitand and plug it in the Array API

Author: agnesLeroy

backends/tfhe-cuda-backend/cuda/src/integer/bitwise_ops.cuh

@@ -25,38 +25,43 @@ __host__ void host_integer_radix_bitop_kb(
           lwe_array_out->num_radix_blocks == lwe_array_2->num_radix_blocks,
       "Cuda error: input and output num radix blocks must be equal");
 
-    PANIC_IF_FALSE(
-            lwe_array_out->num_radix_ciphertexts == lwe_array_1->num_radix_ciphertexts &&
-            lwe_array_out->num_radix_ciphertexts == lwe_array_2->num_radix_ciphertexts,
-            "Cuda error: input and output num radix ciphertexts must be equal");
+  PANIC_IF_FALSE(
+      lwe_array_out->num_radix_ciphertexts ==
+              lwe_array_1->num_radix_ciphertexts &&
+          lwe_array_out->num_radix_ciphertexts ==
+              lwe_array_2->num_radix_ciphertexts,
+      "Cuda error: input and output num radix ciphertexts must be equal");
 
-    PANIC_IF_FALSE(lwe_array_out->lwe_dimension == lwe_array_1->lwe_dimension &&
+  PANIC_IF_FALSE(lwe_array_out->lwe_dimension == lwe_array_1->lwe_dimension &&
                      lwe_array_out->lwe_dimension == lwe_array_2->lwe_dimension,
                  "Cuda error: input and output lwe dimension must be equal");
 
   auto lut = mem_ptr->lut;
-  uint64_t degrees[lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts];
+  uint64_t degrees[lwe_array_1->num_radix_blocks *
+                   lwe_array_1->num_radix_ciphertexts];
   if (mem_ptr->op == BITOP_TYPE::BITAND) {
-    update_degrees_after_bitand(degrees, lwe_array_1->degrees,
-                                lwe_array_2->degrees,
-                                lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts);
+    update_degrees_after_bitand(
+        degrees, lwe_array_1->degrees, lwe_array_2->degrees,
+        lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts);
   } else if (mem_ptr->op == BITOP_TYPE::BITOR) {
-    update_degrees_after_bitor(degrees, lwe_array_1->degrees,
-                               lwe_array_2->degrees,
-                               lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts);
+    update_degrees_after_bitor(
+        degrees, lwe_array_1->degrees, lwe_array_2->degrees,
+        lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts);
   } else if (mem_ptr->op == BITOP_TYPE::BITXOR) {
-    update_degrees_after_bitxor(degrees, lwe_array_1->degrees,
-                                lwe_array_2->degrees,
-                                lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts);
+    update_degrees_after_bitxor(
+        degrees, lwe_array_1->degrees, lwe_array_2->degrees,
+        lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts);
   }
 
   integer_radix_apply_bivariate_lookup_table_kb<Torus>(
       streams, gpu_indexes, gpu_count, lwe_array_out, lwe_array_1, lwe_array_2,
-      bsks, ksks, ms_noise_reduction_key, lut, lwe_array_out->num_radix_blocks * lwe_array_out->num_radix_ciphertexts,
+      bsks, ksks, ms_noise_reduction_key, lut,
+      lwe_array_out->num_radix_blocks * lwe_array_out->num_radix_ciphertexts,
       lut->params.message_modulus);
 
   memcpy(lwe_array_out->degrees, degrees,
-         lwe_array_out->num_radix_blocks * lwe_array_out->num_radix_ciphertexts * sizeof(uint64_t));
+         lwe_array_out->num_radix_blocks *
+             lwe_array_out->num_radix_ciphertexts * sizeof(uint64_t));
 }
 
 template <typename Torus>

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

@@ -19,7 +19,8 @@ use crate::integer::block_decomposition::{
     DecomposableInto, RecomposableFrom, RecomposableSignedInteger,
 };
 use crate::integer::gpu::ciphertext::{
-    CudaIntegerRadixCiphertext, CudaSignedRadixCiphertext, CudaUnsignedRadixCiphertext,
+    CudaIntegerRadixCiphertext, CudaRadixCiphertext, CudaSignedRadixCiphertext,
+    CudaUnsignedRadixCiphertext,
 };
 use crate::integer::server_key::radix_parallel::scalar_div_mod::SignedReciprocable;
 use crate::integer::server_key::{Reciprocable, ScalarMultiplier};
@@ -83,6 +84,12 @@ impl<'a, T> TensorSlice<'a, GpuSlice<'a, T>> {
     pub fn par_iter(self) -> ParStridedIter<'a, T> {
         ParStridedIter::new(self.slice.0, self.dims.clone())
     }
+    pub fn len(&self) -> usize {
+        self.dims.flattened_len()
+    }
+    pub fn as_slice(&self) -> &'a [T] {
+        self.slice.0
+    }
 }
 
 impl<'a, T> TensorSlice<'a, GpuSliceMut<'a, T>> {
@@ -316,7 +323,25 @@ where
         lhs: TensorSlice<'_, Self::Slice<'a>>,
         rhs: TensorSlice<'_, Self::Slice<'a>>,
     ) -> Self::Owned {
-        par_map_sks_op_on_pair_of_elements(lhs, rhs, crate::integer::gpu::CudaServerKey::bitand)
+        GpuOwned(global_state::with_cuda_internal_keys(|cuda_key| {
+            let streams = &cuda_key.streams;
+            let num_ciphertexts = lhs.len() as u32;
+            let lhs_slice: &[T] = lhs.as_slice();
+            let rhs_slice: &[T] = rhs.as_slice();
+            let mut lhs_aligned = T::from(CudaRadixCiphertext::from_radix_ciphertext_vec(
+                lhs_slice, streams,
+            ));
+            let rhs_aligned = T::from(CudaRadixCiphertext::from_radix_ciphertext_vec(
+                rhs_slice, streams,
+            ));
+            crate::integer::gpu::CudaServerKey::bitand_vec(
+                cuda_key.pbs_key(),
+                &mut lhs_aligned,
+                &rhs_aligned,
+                num_ciphertexts,
+                streams,
+            )
+        }))
     }
 
     fn bitor<'a>(

tfhe/src/high_level_api/array/traits.rs

@@ -28,6 +28,12 @@ impl<'a, T> TensorSlice<'a, &'a [T]> {
     pub fn par_iter(self) -> ParStridedIter<'a, T> {
         ParStridedIter::new(self.slice, self.dims.clone())
     }
+    pub fn len(&self) -> usize {
+        self.dims.flattened_len()
+    }
+    pub fn as_slice(&self) -> &'a [T] {
+        self.slice
+    }
 }
 
 impl<'a, T> TensorSlice<'a, &'a mut [T]> {

tfhe/src/integer/gpu/ciphertext/mod.rs

@@ -7,14 +7,15 @@ pub mod squashed_noise;
 
 use crate::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
 use crate::core_crypto::gpu::vec::CudaVec;
-use crate::core_crypto::gpu::CudaStreams;
+use crate::core_crypto::gpu::{CudaLweList, CudaStreams};
 use crate::core_crypto::prelude::{LweCiphertextList, LweCiphertextOwned};
 use crate::integer::gpu::ciphertext::info::{CudaBlockInfo, CudaRadixCiphertextInfo};
 use crate::integer::parameters::LweDimension;
 use crate::integer::{IntegerCiphertext, RadixCiphertext, SignedRadixCiphertext};
 use crate::shortint::{Ciphertext, EncryptionKeyChoice};
 use crate::GpuIndex;
 
+use crate::shortint::parameters::LweCiphertextCount;
 pub use compressed_noise_squashed_ciphertext_list::*;
 
 pub trait CudaIntegerRadixCiphertext: Sized {
@@ -70,6 +71,60 @@ pub trait CudaIntegerRadixCiphertext: Sized {
     fn gpu_indexes(&self) -> &[GpuIndex] {
         &self.as_ref().d_blocks.0.d_vec.gpu_indexes
     }
+
+    fn to_integer_radix_ciphertext_vec(
+        &self,
+        num_radix_ciphertexts: u32,
+        streams: &CudaStreams,
+    ) -> Vec<Self> {
+        let total_blocks = self.as_ref().d_blocks.0.lwe_ciphertext_count.0;
+        assert_eq!(total_blocks % num_radix_ciphertexts as usize, 0, "Total number of blocks ({total_blocks}) is not divisible by number of radix ciphertexts ({num_radix_ciphertexts})");
+
+        let num_blocks = total_blocks / num_radix_ciphertexts as usize;
+
+        let mut result = Vec::with_capacity(num_radix_ciphertexts as usize);
+        let lwe_dimension = self.as_ref().d_blocks.lwe_dimension();
+
+        for i in 0..num_radix_ciphertexts as usize {
+            let block_start = i * num_blocks;
+            let block_end = block_start + num_blocks;
+
+            let d_vec = unsafe {
+                let mut d_vec =
+                    CudaVec::new_async(lwe_dimension.to_lwe_size().0 * num_blocks, streams, 0);
+
+                let copy_start = block_start * lwe_dimension.to_lwe_size().0;
+                let copy_end = (block_end + 1) * lwe_dimension.to_lwe_size().0;
+                d_vec.copy_src_range_gpu_to_gpu_async(
+                    copy_start..copy_end,
+                    &self.as_ref().d_blocks.0.d_vec,
+                    streams,
+                    0,
+                );
+
+                streams.synchronize();
+                d_vec
+            };
+            let lwe_list = CudaLweList::<u64> {
+                d_vec,
+                lwe_ciphertext_count: LweCiphertextCount(num_blocks),
+                lwe_dimension,
+                ciphertext_modulus: self.as_ref().d_blocks.ciphertext_modulus(),
+            };
+
+            // Copy the associated block metadata
+            let block_info = self.as_ref().info.blocks[block_start..block_end].to_vec();
+
+            let info = CudaRadixCiphertextInfo { blocks: block_info };
+
+            result.push(Self::from(CudaRadixCiphertext::new(
+                CudaLweCiphertextList(lwe_list),
+                info,
+            )));
+        }
+
+        result
+    }
 }
 
 pub struct CudaRadixCiphertext {

tfhe/src/integer/gpu/mod.rs

@@ -8134,3 +8134,150 @@ pub unsafe fn expand_async<T: UnsignedInteger, B: Numeric>(
         std::ptr::addr_of_mut!(mem_ptr),
     );
 }
+
+#[allow(clippy::too_many_arguments)]
+/// # Safety
+///
+/// This operation modifies raw GPU pointers on the GPU
+pub unsafe fn unchecked_bitop_vec_radix_kb_assign<T: UnsignedInteger, B: Numeric>(
+    streams: &CudaStreams,
+    radix_lwe_left: &mut CudaRadixCiphertext,
+    radix_lwe_right: &CudaRadixCiphertext,
+    bootstrapping_key: &CudaVec<B>,
+    keyswitch_key: &CudaVec<T>,
+    message_modulus: MessageModulus,
+    carry_modulus: CarryModulus,
+    glwe_dimension: GlweDimension,
+    polynomial_size: PolynomialSize,
+    big_lwe_dimension: LweDimension,
+    small_lwe_dimension: LweDimension,
+    ks_level: DecompositionLevelCount,
+    ks_base_log: DecompositionBaseLog,
+    pbs_level: DecompositionLevelCount,
+    pbs_base_log: DecompositionBaseLog,
+    op: BitOpType,
+    num_blocks: u32,
+    num_radix_ciphertexts: u32,
+    pbs_type: PBSType,
+    grouping_factor: LweBskGroupingFactor,
+    ms_noise_reduction_configuration: Option<&CudaModulusSwitchNoiseReductionConfiguration>,
+) {
+    assert_eq!(
+        streams.gpu_indexes[0],
+        radix_lwe_left.d_blocks.0.d_vec.gpu_index(0),
+        "GPU error: first stream is on GPU {}, first lhs pointer is on GPU {}",
+        streams.gpu_indexes[0].get(),
+        radix_lwe_left.d_blocks.0.d_vec.gpu_index(0).get(),
+    );
+    assert_eq!(
+        streams.gpu_indexes[0],
+        radix_lwe_right.d_blocks.0.d_vec.gpu_index(0),
+        "GPU error: first stream is on GPU {}, first rhs pointer is on GPU {}",
+        streams.gpu_indexes[0].get(),
+        radix_lwe_right.d_blocks.0.d_vec.gpu_index(0).get(),
+    );
+    assert_eq!(
+        streams.gpu_indexes[0],
+        bootstrapping_key.gpu_index(0),
+        "GPU error: first stream is on GPU {}, first bsk pointer is on GPU {}",
+        streams.gpu_indexes[0].get(),
+        bootstrapping_key.gpu_index(0).get(),
+    );
+    assert_eq!(
+        streams.gpu_indexes[0],
+        keyswitch_key.gpu_index(0),
+        "GPU error: first stream is on GPU {}, first ksk pointer is on GPU {}",
+        streams.gpu_indexes[0].get(),
+        keyswitch_key.gpu_index(0).get(),
+    );
+    let ct_modulus = radix_lwe_left
+        .d_blocks
+        .ciphertext_modulus()
+        .raw_modulus_float();
+    let (noise_reduction_type, ms_noise_reduction_key_ffi) =
+        resolve_ms_noise_reduction_config(ms_noise_reduction_configuration, ct_modulus);
+
+    let mut mem_ptr: *mut i8 = std::ptr::null_mut();
+    let mut radix_lwe_left_degrees = radix_lwe_left
+        .info
+        .blocks
+        .iter()
+        .map(|b| b.degree.0)
+        .collect();
+    let mut radix_lwe_left_noise_levels = radix_lwe_left
+        .info
+        .blocks
+        .iter()
+        .map(|b| b.noise_level.0)
+        .collect();
+    let mut cuda_ffi_radix_lwe_left = prepare_cuda_radix_ffi(
+        radix_lwe_left,
+        &mut radix_lwe_left_degrees,
+        &mut radix_lwe_left_noise_levels,
+    );
+    // Here even though the input is not modified, data is passed as mutable.
+    // This avoids having to create two structs for the CudaRadixCiphertext pointers,
+    // one const and the other mutable.
+    // Having two structs on the Cuda side complicates things as we need to be sure we pass the
+    // Const structure as input instead of the mutable structure, which leads to complicated
+    // data manipulation on the C++ side to change mutability of data.
+    let mut radix_lwe_right_degrees = radix_lwe_right
+        .info
+        .blocks
+        .iter()
+        .map(|b| b.degree.0)
+        .collect();
+    let mut radix_lwe_right_noise_levels = radix_lwe_right
+        .info
+        .blocks
+        .iter()
+        .map(|b| b.noise_level.0)
+        .collect();
+    let cuda_ffi_radix_lwe_right = prepare_cuda_radix_ffi(
+        radix_lwe_right,
+        &mut radix_lwe_right_degrees,
+        &mut radix_lwe_right_noise_levels,
+    );
+    scratch_cuda_integer_radix_bitop_kb_64(
+        streams.ptr.as_ptr(),
+        streams.gpu_indexes_ptr(),
+        streams.len() as u32,
+        std::ptr::addr_of_mut!(mem_ptr),
+        glwe_dimension.0 as u32,
+        polynomial_size.0 as u32,
+        big_lwe_dimension.0 as u32,
+        small_lwe_dimension.0 as u32,
+        ks_level.0 as u32,
+        ks_base_log.0 as u32,
+        pbs_level.0 as u32,
+        pbs_base_log.0 as u32,
+        grouping_factor.0 as u32,
+        num_blocks * num_radix_ciphertexts,
+        message_modulus.0 as u32,
+        carry_modulus.0 as u32,
+        pbs_type as u32,
+        op as u32,
+        true,
+        noise_reduction_type as u32,
+    );
+    cuda_bitop_integer_radix_ciphertext_kb_64(
+        streams.ptr.as_ptr(),
+        streams.gpu_indexes_ptr(),
+        streams.len() as u32,
+        &raw mut cuda_ffi_radix_lwe_left,
+        &raw const cuda_ffi_radix_lwe_left,
+        &raw const cuda_ffi_radix_lwe_right,
+        mem_ptr,
+        bootstrapping_key.ptr.as_ptr(),
+        keyswitch_key.ptr.as_ptr(),
+        &raw const ms_noise_reduction_key_ffi,
+    );
+    cleanup_cuda_integer_bitop(
+        streams.ptr.as_ptr(),
+        streams.gpu_indexes_ptr(),
+        streams.len() as u32,
+        std::ptr::addr_of_mut!(mem_ptr),
+    );
+    update_noise_degree(radix_lwe_left, &cuda_ffi_radix_lwe_left);
+    streams.synchronize();
+}