tfhe/high_level_api/
gpu_utils.rs

1use crate::GpuIndex;
2
3/// Check there is enough memory on the GPU to allocate a certain size
4/// # Example
5///
6/// ```rust
7/// use rand::Rng;
8/// use tfhe::prelude::*;
9/// use tfhe::{
10///     generate_keys, set_server_key, ClientKey, CompressedServerKey, ConfigBuilder, FheInt16,
11///     FheInt32, GpuIndex,
12/// };
13///
14/// let config = ConfigBuilder::default();
15/// let client_key = ClientKey::generate(config);
16/// let csks = CompressedServerKey::new(&client_key);
17/// let server_key = csks.decompress_to_gpu();
18/// set_server_key(server_key);
19/// let mut rng = rand::thread_rng();
20/// let clear_a = rng.gen_range(1..=i32::MAX);
21/// let clear_b = rng.gen_range(1..=i32::MAX);
22/// let mut a = FheInt32::try_encrypt(clear_a, &client_key).unwrap();
23/// let mut b = FheInt32::try_encrypt(clear_b, &client_key).unwrap();
24/// let ciphertexts_size = a.get_size_on_gpu() + b.get_size_on_gpu();
25/// assert!(check_valid_cuda_malloc(ciphertexts_size, GpuIndex::new(0)));
26///
27/// a.move_to_current_device();
28/// b.move_to_current_device();
29///
30/// let tmp_buffer_size = a.get_add_size_on_gpu(&b);
31/// assert!(check_valid_cuda_malloc(tmp_buffer_size, GpuIndex::new(0)));
32/// a += &b;
33/// ```
34pub fn check_valid_cuda_malloc(size: u64, gpu_index: GpuIndex) -> bool {
35    crate::core_crypto::gpu::check_valid_cuda_malloc(size, gpu_index)
36}
tfhe/high_level_api/gpu_utils.rs

tfhe/high_level_api/
gpu_utils.rs
