chore(gpu): reintroduce a GpuIndex type

zama-ai · Nov 29, 2024 · df68e54 · df68e54
1 parent 2bf19b6
commit df68e54
Show file tree

Hide file tree

Showing 34 changed files with 1,654 additions and 1,110 deletions.
diff --git a/tfhe/benches/core_crypto/ks_bench.rs b/tfhe/benches/core_crypto/ks_bench.rs
@@ -307,7 +307,7 @@ mod cuda {
     use tfhe::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
     use tfhe::core_crypto::gpu::lwe_keyswitch_key::CudaLweKeyswitchKey;
     use tfhe::core_crypto::gpu::lwe_packing_keyswitch_key::CudaLwePackingKeyswitchKey;
-    use tfhe::core_crypto::gpu::vec::CudaVec;
+    use tfhe::core_crypto::gpu::vec::{CudaVec, GpuIndex};
     use tfhe::core_crypto::gpu::{
         cuda_keyswitch_lwe_ciphertext, cuda_keyswitch_lwe_ciphertext_list_into_glwe_ciphertext,
         CudaStreams,
@@ -330,7 +330,7 @@ mod cuda {
             SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());
 
         let gpu_index = 0;
-        let streams = CudaStreams::new_single_gpu(gpu_index);
+        let streams = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
         for (name, params) in parameters.iter() {
             let lwe_dimension = params.lwe_dimension.unwrap();
@@ -433,7 +433,7 @@ mod cuda {
             SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());
 
         let gpu_index = 0;
-        let streams = CudaStreams::new_single_gpu(gpu_index);
+        let streams = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
         for (name, params) in parameters.iter() {
             let lwe_dimension = params.lwe_dimension.unwrap();

diff --git a/tfhe/benches/core_crypto/pbs_bench.rs b/tfhe/benches/core_crypto/pbs_bench.rs
@@ -830,7 +830,7 @@ mod cuda {
     use tfhe::core_crypto::gpu::lwe_bootstrap_key::CudaLweBootstrapKey;
     use tfhe::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
     use tfhe::core_crypto::gpu::lwe_multi_bit_bootstrap_key::CudaLweMultiBitBootstrapKey;
-    use tfhe::core_crypto::gpu::vec::CudaVec;
+    use tfhe::core_crypto::gpu::vec::{CudaVec, GpuIndex};
     use tfhe::core_crypto::gpu::{
         cuda_multi_bit_programmable_bootstrap_lwe_ciphertext,
         cuda_programmable_bootstrap_lwe_ciphertext, CudaStreams,
@@ -900,7 +900,7 @@ mod cuda {
             SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());
 
         let gpu_index = 0;
-        let stream = CudaStreams::new_single_gpu(gpu_index);
+        let stream = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
         for (name, params) in parameters.iter() {
             // Create the LweSecretKey
@@ -1021,7 +1021,7 @@ mod cuda {
             SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());
 
         let gpu_index = 0;
-        let stream = CudaStreams::new_single_gpu(gpu_index);
+        let stream = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
         for (name, params, grouping_factor) in parameters.iter() {
             // Create the LweSecretKey
@@ -1143,7 +1143,7 @@ mod cuda {
             SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());
 
         let gpu_index = 0;
-        let stream = CudaStreams::new_single_gpu(gpu_index);
+        let stream = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
         for (name, params) in parameters.iter() {
             let input_lwe_secret_key = allocate_and_generate_new_binary_lwe_secret_key(
@@ -1283,7 +1283,7 @@ mod cuda {
             SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());
 
         let gpu_index = 0;
-        let stream = CudaStreams::new_single_gpu(gpu_index);
+        let stream = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
         for (name, params, grouping_factor) in parameters.iter() {
             let input_lwe_secret_key = allocate_and_generate_new_binary_lwe_secret_key(

diff --git a/tfhe/src/core_crypto/gpu/algorithms/glwe_sample_extraction.rs b/tfhe/src/core_crypto/gpu/algorithms/glwe_sample_extraction.rs
@@ -46,7 +46,7 @@ pub unsafe fn cuda_extract_lwe_samples_from_glwe_ciphertext_list_async<Scalar>(
     let nth_array: Vec<u32> = vec_nth.iter().map(|x| x.0 as u32).collect_vec();
     let gpu_indexes = &streams.gpu_indexes;
     unsafe {
-        let d_nth_array = CudaVec::from_cpu_async(&nth_array, streams, gpu_indexes[0]);
+        let d_nth_array = CudaVec::from_cpu_async(&nth_array, streams, gpu_indexes[0].0);
         extract_lwe_samples_from_glwe_ciphertext_list_async(
             streams,
             &mut output_lwe_list.0.d_vec,

diff --git a/tfhe/src/core_crypto/gpu/algorithms/test/glwe_sample_extraction.rs b/tfhe/src/core_crypto/gpu/algorithms/test/glwe_sample_extraction.rs
@@ -2,6 +2,7 @@ use super::*;
 use crate::core_crypto::gpu::glwe_ciphertext_list::CudaGlweCiphertextList;
 use crate::core_crypto::gpu::glwe_sample_extraction::cuda_extract_lwe_samples_from_glwe_ciphertext_list;
 use crate::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
+use crate::core_crypto::gpu::vec::GpuIndex;
 use crate::core_crypto::gpu::CudaStreams;
 use itertools::Itertools;
 
@@ -26,7 +27,7 @@ fn glwe_encrypt_sample_extract_decrypt_custom_mod<Scalar: UnsignedTorus + Send +
     let delta: Scalar = encoding_with_padding / msg_modulus;
 
     let gpu_index = 0;
-    let streams = CudaStreams::new_single_gpu(gpu_index);
+    let streams = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
     let mut msgs = vec![];
 

diff --git a/tfhe/src/core_crypto/gpu/algorithms/test/lwe_keyswitch.rs b/tfhe/src/core_crypto/gpu/algorithms/test/lwe_keyswitch.rs
@@ -1,7 +1,7 @@
 use super::*;
 use crate::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
 use crate::core_crypto::gpu::lwe_keyswitch_key::CudaLweKeyswitchKey;
-use crate::core_crypto::gpu::vec::CudaVec;
+use crate::core_crypto::gpu::vec::{CudaVec, GpuIndex};
 use crate::core_crypto::gpu::{cuda_keyswitch_lwe_ciphertext, CudaStreams};
 use itertools::Itertools;
 
@@ -18,7 +18,7 @@ fn lwe_encrypt_ks_decrypt_custom_mod<Scalar: UnsignedTorus + CastFrom<usize>>(
     let ks_decomp_base_log = params.ks_base_log;
     let ks_decomp_level_count = params.ks_level;
 
-    let stream = CudaStreams::new_single_gpu(0);
+    let stream = CudaStreams::new_single_gpu(GpuIndex(0));
 
     let mut rsc = TestResources::new();
 

diff --git a/tfhe/src/core_crypto/gpu/algorithms/test/lwe_multi_bit_programmable_bootstrapping.rs b/tfhe/src/core_crypto/gpu/algorithms/test/lwe_multi_bit_programmable_bootstrapping.rs
@@ -2,7 +2,7 @@ use super::*;
 use crate::core_crypto::gpu::glwe_ciphertext_list::CudaGlweCiphertextList;
 use crate::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
 use crate::core_crypto::gpu::lwe_multi_bit_bootstrap_key::CudaLweMultiBitBootstrapKey;
-use crate::core_crypto::gpu::vec::CudaVec;
+use crate::core_crypto::gpu::vec::{CudaVec, GpuIndex};
 use crate::core_crypto::gpu::{cuda_multi_bit_programmable_bootstrap_lwe_ciphertext, CudaStreams};
 use itertools::Itertools;
 
@@ -28,7 +28,7 @@ fn lwe_encrypt_multi_bit_pbs_decrypt_custom_mod<
     let grouping_factor = params.grouping_factor;
 
     let gpu_index = 0;
-    let stream = CudaStreams::new_single_gpu(gpu_index);
+    let stream = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
     let mut rsc = TestResources::new();
 

diff --git a/tfhe/src/core_crypto/gpu/algorithms/test/lwe_packing_keyswitch.rs b/tfhe/src/core_crypto/gpu/algorithms/test/lwe_packing_keyswitch.rs
@@ -2,6 +2,7 @@ use super::*;
 use crate::core_crypto::gpu::algorithms::lwe_packing_keyswitch::cuda_keyswitch_lwe_ciphertext_list_into_glwe_ciphertext_async;
 use crate::core_crypto::gpu::glwe_ciphertext_list::CudaGlweCiphertextList;
 use crate::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
+use crate::core_crypto::gpu::vec::GpuIndex;
 use crate::core_crypto::gpu::CudaStreams;
 use serde::de::DeserializeOwned;
 use serde::Serialize;
@@ -67,7 +68,7 @@ where
     let delta: Scalar = encoding_with_padding / msg_modulus;
 
     let gpu_index = 0;
-    let stream = CudaStreams::new_single_gpu(gpu_index);
+    let stream = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
     while msg != Scalar::ZERO {
         msg = msg.wrapping_sub(Scalar::ONE);
@@ -151,7 +152,7 @@ where
     let delta: Scalar = encoding_with_padding / msg_modulus;
 
     let gpu_index = 0;
-    let stream = CudaStreams::new_single_gpu(gpu_index);
+    let stream = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
     while msg != Scalar::ZERO {
         msg = msg.wrapping_sub(Scalar::ONE);

diff --git a/tfhe/src/core_crypto/gpu/algorithms/test/lwe_programmable_bootstrapping.rs b/tfhe/src/core_crypto/gpu/algorithms/test/lwe_programmable_bootstrapping.rs
@@ -2,7 +2,7 @@ use super::*;
 use crate::core_crypto::gpu::glwe_ciphertext_list::CudaGlweCiphertextList;
 use crate::core_crypto::gpu::lwe_bootstrap_key::CudaLweBootstrapKey;
 use crate::core_crypto::gpu::lwe_ciphertext_list::CudaLweCiphertextList;
-use crate::core_crypto::gpu::vec::CudaVec;
+use crate::core_crypto::gpu::vec::{CudaVec, GpuIndex};
 use crate::core_crypto::gpu::{cuda_programmable_bootstrap_lwe_ciphertext, CudaStreams};
 use itertools::Itertools;
 
@@ -26,7 +26,7 @@ fn lwe_encrypt_pbs_decrypt<
     let decomp_level_count = params.pbs_level;
 
     let gpu_index = 0;
-    let stream = CudaStreams::new_single_gpu(gpu_index);
+    let stream = CudaStreams::new_single_gpu(GpuIndex(gpu_index));
 
     let mut rsc = TestResources::new();
 

diff --git a/tfhe/src/core_crypto/gpu/entities/lwe_ciphertext_list.rs b/tfhe/src/core_crypto/gpu/entities/lwe_ciphertext_list.rs
@@ -128,7 +128,7 @@ impl<T: UnsignedInteger> CudaLweCiphertextList<T> {
                 first_item.0.d_vec.as_c_ptr(0),
                 size as u64,
                 streams.ptr[0],
-                streams.gpu_indexes[0],
+                streams.gpu_indexes[0].0,
             );
             ptr = ptr.wrapping_byte_add(size);
             for list in cuda_ciphertexts_list_vec {
@@ -137,7 +137,7 @@ impl<T: UnsignedInteger> CudaLweCiphertextList<T> {
                     list.0.d_vec.as_c_ptr(0),
                     size as u64,
                     streams.ptr[0],
-                    streams.gpu_indexes[0],
+                    streams.gpu_indexes[0].0,
                 );
                 ptr = ptr.wrapping_byte_add(size);
             }