feat(gpu): add new tests to cover the multi-device execution of integ…

…er operations - multiple operations are issued in parallel running independently on different devices, - tests will only run when more than 1 GPU is available, - we only test ERC20-related operators: (overflow_) add/sub, cmp, and if_then_else.
zama-ai · Jan 31, 2025 · cc6edd0 · cc6edd0
1 parent 64b8d1e
commit cc6edd0
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 1 deletion.
diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_unsigned/test_add.rs b/tfhe/src/integer/gpu/server_key/radix/tests_unsigned/test_add.rs
@@ -1,5 +1,6 @@
-use crate::core_crypto::gpu::CudaStreams;
+use crate::core_crypto::gpu::{get_number_of_gpus, CudaStreams};
 use crate::integer::gpu::ciphertext::CudaUnsignedRadixCiphertext;
+use crate::integer::gpu::server_key::radix::tests_long_run::GpuMultiDeviceFunctionExecutor;
 use crate::integer::gpu::server_key::radix::tests_unsigned::{
     create_gpu_parameterized_test, GpuFunctionExecutor,
 };
@@ -14,8 +15,10 @@ use crate::shortint::parameters::*;
 create_gpu_parameterized_test!(integer_unchecked_add);
 create_gpu_parameterized_test!(integer_unchecked_add_assign);
 create_gpu_parameterized_test!(integer_add);
+create_gpu_parameterized_test!(multi_device_integer_add);
 create_gpu_parameterized_test!(integer_sum_ciphertexts_vec);
 create_gpu_parameterized_test!(integer_default_overflowing_add);
+create_gpu_parameterized_test!(multi_device_integer_default_overflowing_add);
 
 fn integer_unchecked_add<P>(param: P)
 where
@@ -41,6 +44,17 @@ where
     default_add_test(param, executor);
 }
 
+fn multi_device_integer_add<P>(param: P)
+where
+    P: Into<PBSParameters>,
+{
+    let executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::add);
+    let num_gpus = get_number_of_gpus();
+    if num_gpus > 1 {
+        default_add_test(param, executor);
+    }
+}
+
 fn integer_sum_ciphertexts_vec<P>(param: P)
 where
     P: Into<PBSParameters>,
@@ -65,3 +79,14 @@ where
     let executor = GpuFunctionExecutor::new(&CudaServerKey::unsigned_overflowing_add);
     default_overflowing_add_test(param, executor);
 }
+
+fn multi_device_integer_default_overflowing_add<P>(param: P)
+where
+    P: Into<PBSParameters>,
+{
+    let executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::unsigned_overflowing_add);
+    let num_gpus = get_number_of_gpus();
+    if num_gpus > 1 {
+        default_overflowing_add_test(param, executor);
+    }
+}
diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_unsigned/test_cmux.rs b/tfhe/src/integer/gpu/server_key/radix/tests_unsigned/test_cmux.rs
@@ -1,3 +1,5 @@
+use crate::core_crypto::gpu::get_number_of_gpus;
+use crate::integer::gpu::server_key::radix::tests_long_run::GpuMultiDeviceFunctionExecutor;
 use crate::integer::gpu::server_key::radix::tests_unsigned::{
     create_gpu_parameterized_test, GpuFunctionExecutor,
 };
@@ -6,6 +8,7 @@ use crate::integer::server_key::radix_parallel::tests_unsigned::test_cmux::defau
 use crate::shortint::parameters::*;
 
 create_gpu_parameterized_test!(integer_if_then_else);
+create_gpu_parameterized_test!(multi_device_integer_if_then_else);
 
 fn integer_if_then_else<P>(param: P)
 where
@@ -14,3 +17,13 @@ where
     let executor = GpuFunctionExecutor::new(&CudaServerKey::if_then_else);
     default_if_then_else_test(param, executor);
 }
+fn multi_device_integer_if_then_else<P>(param: P)
+where
+    P: Into<PBSParameters>,
+{
+    let executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::if_then_else);
+    let num_gpus = get_number_of_gpus();
+    if num_gpus > 1 {
+        default_if_then_else_test(param, executor);
+    }
+}
diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_unsigned/test_comparison.rs b/tfhe/src/integer/gpu/server_key/radix/tests_unsigned/test_comparison.rs
@@ -1,3 +1,5 @@
+use crate::core_crypto::gpu::get_number_of_gpus;
+use crate::integer::gpu::server_key::radix::tests_long_run::GpuMultiDeviceFunctionExecutor;
 use crate::integer::gpu::server_key::radix::tests_unsigned::{
     create_gpu_parameterized_test, GpuFunctionExecutor,
 };
@@ -40,6 +42,20 @@ macro_rules! define_gpu_comparison_test_functions {
                 )
             }
 
+            fn [<multi_device_integer_default_ $comparison_name _ $clear_type:lower>]<P>(param: P) where P: Into<PBSParameters> {
+                let num_tests = 1;
+                let executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::[<$comparison_name>]);
+                let num_gpus = get_number_of_gpus();
+                if num_gpus > 1 {
+                    test_default_function(
+                        param,
+                        num_tests,
+                        executor,
+                        |lhs, rhs| $clear_type::from(<$clear_type>::$comparison_name(&lhs, &rhs)),
+                    )
+                }
+            }
+
             create_gpu_parameterized_test!([<integer_unchecked_ $comparison_name _ $clear_type:lower>]{
                 PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
                 PARAM_GPU_MULTI_BIT_GROUP_3_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
@@ -50,6 +66,11 @@ macro_rules! define_gpu_comparison_test_functions {
                 PARAM_GPU_MULTI_BIT_GROUP_3_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
                 V1_0_PARAM_GPU_MULTI_BIT_GROUP_2_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64,
             });
+            create_gpu_parameterized_test!([<multi_device_integer_default_ $comparison_name _ $clear_type:lower>]{
+                PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
+                PARAM_GPU_MULTI_BIT_GROUP_3_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
+                V1_0_PARAM_GPU_MULTI_BIT_GROUP_2_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64,
+            });
         }
     };
 }

diff --git a/tfhe/src/integer/gpu/server_key/radix/tests_unsigned/test_sub.rs b/tfhe/src/integer/gpu/server_key/radix/tests_unsigned/test_sub.rs
@@ -1,3 +1,5 @@
+use crate::core_crypto::gpu::get_number_of_gpus;
+use crate::integer::gpu::server_key::radix::tests_long_run::GpuMultiDeviceFunctionExecutor;
 use crate::integer::gpu::server_key::radix::tests_unsigned::{
     create_gpu_parameterized_test, GpuFunctionExecutor,
 };
@@ -10,7 +12,9 @@ use crate::shortint::parameters::*;
 
 create_gpu_parameterized_test!(integer_unchecked_sub);
 create_gpu_parameterized_test!(integer_sub);
+create_gpu_parameterized_test!(multi_device_integer_sub);
 create_gpu_parameterized_test!(integer_default_overflowing_sub);
+create_gpu_parameterized_test!(multi_device_integer_default_overflowing_sub);
 
 fn integer_unchecked_sub<P>(param: P)
 where
@@ -28,10 +32,32 @@ where
     default_sub_test(param, executor);
 }
 
+fn multi_device_integer_sub<P>(param: P)
+where
+    P: Into<PBSParameters>,
+{
+    let executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::sub);
+    let num_gpus = get_number_of_gpus();
+    if num_gpus > 1 {
+        default_sub_test(param, executor);
+    }
+}
+
 fn integer_default_overflowing_sub<P>(param: P)
 where
     P: Into<PBSParameters>,
 {
     let executor = GpuFunctionExecutor::new(&CudaServerKey::unsigned_overflowing_sub);
     default_overflowing_sub_test(param, executor);
 }
+
+fn multi_device_integer_default_overflowing_sub<P>(param: P)
+where
+    P: Into<PBSParameters>,
+{
+    let executor = GpuMultiDeviceFunctionExecutor::new(&CudaServerKey::unsigned_overflowing_sub);
+    let num_gpus = get_number_of_gpus();
+    if num_gpus > 1 {
+        default_overflowing_sub_test(param, executor);
+    }
+}