Skip to content

Commit

Permalink
chore(ci): add randomized long run tests on CPU and GPU
Browse files Browse the repository at this point in the history
  • Loading branch information
agnesLeroy committed Dec 5, 2024
1 parent fd31694 commit 7abc896
Show file tree
Hide file tree
Showing 6 changed files with 1,732 additions and 2 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/gpu_integer_long_run_tests.yml
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
name: AWS Long Run Tests on GPU
name: Long Run Tests on GPU

env:
CARGO_TERM_COLOR: always
Expand Down
6 changes: 5 additions & 1 deletion tfhe/src/integer/gpu/server_key/radix/tests_long_run/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,13 @@ use crate::integer::gpu::server_key::radix::tests_unsigned::GpuContext;
use crate::integer::gpu::CudaServerKey;
use crate::integer::server_key::radix_parallel::tests_cases_unsigned::FunctionExecutor;
use crate::integer::{BooleanBlock, RadixCiphertext, RadixClientKey, ServerKey, U256};
use rand::Rng;
use std::sync::Arc;
use tfhe_cuda_backend::cuda_bind::cuda_get_number_of_gpus;

pub(crate) mod test_erc20;
pub(crate) mod test_random_op_sequence;

pub(crate) struct GpuMultiDeviceFunctionExecutor<F> {
pub(crate) context: Option<GpuContext>,
pub(crate) func: F,
Expand All @@ -27,7 +30,8 @@ impl<F> GpuMultiDeviceFunctionExecutor<F> {
impl<F> GpuMultiDeviceFunctionExecutor<F> {
pub(crate) fn setup_from_keys(&mut self, cks: &RadixClientKey, _sks: &Arc<ServerKey>) {
let num_gpus = unsafe { cuda_get_number_of_gpus() } as u32;
let streams = CudaStreams::new_single_gpu(GpuIndex(num_gpus - 1));
let gpu_index = GpuIndex(rand::thread_rng().gen_range(0..num_gpus));
let streams = CudaStreams::new_single_gpu(gpu_index);

let sks = CudaServerKey::new(cks.as_ref(), &streams);
streams.synchronize();
Expand Down
Loading

0 comments on commit 7abc896

Please sign in to comment.