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

Author: agnesLeroy

backends/tfhe-cuda-backend/cuda/include/integer/integer.h

@@ -56,6 +56,7 @@ typedef struct {
   uint32_t num_radix_blocks;
   uint32_t max_num_radix_blocks;
   uint32_t lwe_dimension;
+  uint32_t num_radix_ciphertexts;
 } CudaRadixCiphertextFFI;
 
 typedef struct {

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

@@ -24,33 +24,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->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];
+  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);
+    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);
+    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);
+    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, lwe_array_out, lwe_array_1, lwe_array_2, bsks, ksks,
-      ms_noise_reduction_key, lut, lwe_array_out->num_radix_blocks,
+      streams, 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,
       lut->params.message_modulus);
 
   memcpy(lwe_array_out->degrees, degrees,
-         lwe_array_out->num_radix_blocks * sizeof(uint64_t));
+         lwe_array_out->num_radix_blocks *
+             lwe_array_out->num_radix_ciphertexts * sizeof(uint64_t));
 }
 
 template <typename Torus>

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

@@ -417,9 +417,12 @@ __host__ void host_pack_bivariate_blocks(
       lwe_array_out->lwe_dimension != lwe_array_2->lwe_dimension)
     PANIC("Cuda error: input and output radix ciphertexts should have the same "
           "lwe dimension")
-  if (num_radix_blocks > lwe_array_out->num_radix_blocks ||
-      num_radix_blocks > lwe_array_1->num_radix_blocks ||
-      num_radix_blocks > lwe_array_2->num_radix_blocks)
+  if (num_radix_blocks > lwe_array_out->num_radix_blocks *
+                             lwe_array_out->num_radix_ciphertexts ||
+      num_radix_blocks >
+          lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts ||
+      num_radix_blocks >
+          lwe_array_2->num_radix_blocks * lwe_array_2->num_radix_ciphertexts)
     PANIC("Cuda error: num radix blocks on which packing is applied should be "
           "smaller or equal to the number of input & output radix blocks")
 
@@ -530,7 +533,8 @@ __host__ void integer_radix_apply_univariate_lookup_table_kb(
   if (num_radix_blocks > lut->num_blocks)
     PANIC("Cuda error: num radix blocks on which lut is applied should be "
           "smaller or equal to the number of lut radix blocks")
-  if (num_radix_blocks > lwe_array_out->num_radix_blocks)
+  if (num_radix_blocks >
+      lwe_array_out->num_radix_blocks * lwe_array_out->num_radix_ciphertexts)
     PANIC("Cuda error: num radix blocks on which lut is applied should be "
           "smaller or equal to the number of input & output radix blocks")
 
@@ -756,11 +760,14 @@ __host__ void integer_radix_apply_bivariate_lookup_table_kb(
   if (num_radix_blocks > lut->num_blocks)
     PANIC("Cuda error: num radix blocks on which lut is applied should be "
           "smaller or equal to the number of lut radix blocks")
-  if (num_radix_blocks > lwe_array_out->num_radix_blocks ||
-      num_radix_blocks > lwe_array_1->num_radix_blocks ||
-      num_radix_blocks > lwe_array_2->num_radix_blocks)
+  if (num_radix_blocks > lwe_array_out->num_radix_blocks *
+                             lwe_array_out->num_radix_ciphertexts ||
+      num_radix_blocks >
+          lwe_array_1->num_radix_blocks * lwe_array_1->num_radix_ciphertexts ||
+      num_radix_blocks >
+          lwe_array_2->num_radix_blocks * lwe_array_2->num_radix_ciphertexts)
     PANIC("Cuda error: num radix blocks on which lut is applied should be "
-          "smaller or equal to the number of input & output radix blocks")
+          "smaller or equal to the number of total input & output radix blocks")
 
   auto params = lut->params;
   auto pbs_type = params.pbs_type;

backends/tfhe-cuda-backend/cuda/src/integer/radix_ciphertext.cu

@@ -25,6 +25,7 @@ void into_radix_ciphertext(CudaRadixCiphertextFFI *radix, void *lwe_array,
   radix->num_radix_blocks = num_radix_blocks;
   radix->max_num_radix_blocks = num_radix_blocks;
   radix->ptr = lwe_array;
+  radix->num_radix_ciphertexts = 1;
 
   radix->degrees = (uint64_t *)(calloc(num_radix_blocks, sizeof(uint64_t)));
   radix->noise_levels =

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

@@ -19,6 +19,7 @@ void create_zero_radix_ciphertext_async(cudaStream_t const stream,
   radix->lwe_dimension = lwe_dimension;
   radix->num_radix_blocks = num_radix_blocks;
   radix->max_num_radix_blocks = num_radix_blocks;
+  radix->num_radix_ciphertexts = 1;
   uint64_t size = (lwe_dimension + 1) * num_radix_blocks * sizeof(Torus);
   radix->ptr = (void *)cuda_malloc_with_size_tracking_async(
       size, stream, gpu_index, size_tracker, allocate_gpu_memory);
@@ -63,6 +64,7 @@ void as_radix_ciphertext_slice(CudaRadixCiphertextFFI *output_radix,
 
   auto lwe_size = input_radix->lwe_dimension + 1;
   output_radix->num_radix_blocks = end_input_lwe_index - start_input_lwe_index;
+  output_radix->num_radix_ciphertexts = input_radix->num_radix_ciphertexts;
   output_radix->max_num_radix_blocks = input_radix->max_num_radix_blocks;
   output_radix->lwe_dimension = input_radix->lwe_dimension;
   Torus *in_ptr = (Torus *)input_radix->ptr;

backends/tfhe-cuda-backend/src/bindings.rs

@@ -183,6 +183,7 @@ pub struct CudaRadixCiphertextFFI {
     pub num_radix_blocks: u32,
     pub max_num_radix_blocks: u32,
     pub lwe_dimension: u32,
+    pub num_radix_ciphertexts: u32,
 }
 #[allow(clippy::unnecessary_operation, clippy::identity_op)]
 const _: () = {
@@ -201,6 +202,8 @@ const _: () = {
         [::std::mem::offset_of!(CudaRadixCiphertextFFI, max_num_radix_blocks) - 28usize];
     ["Offset of field: CudaRadixCiphertextFFI::lwe_dimension"]
         [::std::mem::offset_of!(CudaRadixCiphertextFFI, lwe_dimension) - 32usize];
+    ["Offset of field: CudaRadixCiphertextFFI::num_radix_ciphertexts"]
+        [::std::mem::offset_of!(CudaRadixCiphertextFFI, num_radix_ciphertexts) - 36usize];
 };
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]

tfhe/src/core_crypto/gpu/mod.rs

@@ -840,6 +840,7 @@ pub unsafe fn add_lwe_ciphertext_vector_async<T: UnsignedInteger>(
         num_radix_blocks: num_samples,
         max_num_radix_blocks: num_samples,
         lwe_dimension: lwe_dimension.0 as u32,
+        num_radix_ciphertexts: 1u32,
     };
     let lwe_array_in_1_data = CudaRadixCiphertextFFI {
         ptr: lwe_array_in_1.get_mut_c_ptr(0),
@@ -848,6 +849,7 @@ pub unsafe fn add_lwe_ciphertext_vector_async<T: UnsignedInteger>(
         num_radix_blocks: num_samples,
         max_num_radix_blocks: num_samples,
         lwe_dimension: lwe_dimension.0 as u32,
+        num_radix_ciphertexts: 1u32,
     };
     let lwe_array_in_2_data = CudaRadixCiphertextFFI {
         ptr: lwe_array_in_2.get_mut_c_ptr(0),
@@ -856,6 +858,7 @@ pub unsafe fn add_lwe_ciphertext_vector_async<T: UnsignedInteger>(
         num_radix_blocks: num_samples,
         max_num_radix_blocks: num_samples,
         lwe_dimension: lwe_dimension.0 as u32,
+        num_radix_ciphertexts: 1u32,
     };
     cuda_add_lwe_ciphertext_vector_64(
         streams.ptr[0],
@@ -890,6 +893,7 @@ pub unsafe fn add_lwe_ciphertext_vector_assign_async<T: UnsignedInteger>(
         num_radix_blocks: num_samples,
         max_num_radix_blocks: num_samples,
         lwe_dimension: lwe_dimension.0 as u32,
+        num_radix_ciphertexts: 1u32,
     };
     let lwe_array_in_data = CudaRadixCiphertextFFI {
         ptr: lwe_array_in.get_mut_c_ptr(0),
@@ -898,6 +902,7 @@ pub unsafe fn add_lwe_ciphertext_vector_assign_async<T: UnsignedInteger>(
         num_radix_blocks: num_samples,
         max_num_radix_blocks: num_samples,
         lwe_dimension: lwe_dimension.0 as u32,
+        num_radix_ciphertexts: 1u32,
     };
     cuda_add_lwe_ciphertext_vector_64(
         streams.ptr[0],

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,8 +71,68 @@ pub trait CudaIntegerRadixCiphertext: Sized {
     fn gpu_indexes(&self) -> &[GpuIndex] {
         &self.as_ref().d_blocks.0.d_vec.gpu_indexes
     }
+
+    // Converts a CudaIntegerRadixCiphertext with num_blocks * num_ciphertexts LWEs into a
+    // Vec<CudaIntegerRadixCiphertext> of length num_radix_ciphertexts, where each ciphertext has
+    // num_blocks LWEs
+    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 * 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
+    }
 }
 
+/// This struct corresponds to the pointers on GPU and
+/// metadata representing an array of LWEs corresponding
+/// to one or more RadixCiphertexts
 pub struct CudaRadixCiphertext {
     pub d_blocks: CudaLweCiphertextList<u64>,
     pub info: CudaRadixCiphertextInfo,