github/tfhe-rs
1
1
Fork 0
mirror of https://github.com/zama-ai/tfhe-rs.git synced 2026-01-07 22:04:10 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix(gpu): fix memory leak in rerand

Browse Source
This commit is contained in:
Agnes Leroy
2025-11-14 11:07:27 +01:00
committed by Agnès Leroy
parent d75844dea5
commit 4f9f4982f6
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
backends/tfhe-cuda-backend/cuda/include/integer/rerand_utilities.h
View File

@@ -7,10 +7,10 @@
template <typename Torus> struct int_rerand_mem {
int_radix_params params;
Torus *lwe_trivial_indexes;
Torus *tmp_zero_lwes;
Torus *tmp_ksed_zero_lwes;
Torus *lwe_trivial_indexes;
uint32_t num_lwes;
bool gpu_memory_allocated;
@@ -54,7 +54,7 @@ template <typename Torus> struct int_rerand_mem {
num_lwes * sizeof(Torus), streams.stream(0),
streams.gpu_index(0));
streams.synchronize();
cuda_synchronize_stream(streams.stream(0), streams.gpu_index(0));
free(h_lwe_trivial_indexes);
}
@@ -66,6 +66,9 @@ template <typename Torus> struct int_rerand_mem {
cuda_drop_with_size_tracking_async(tmp_ksed_zero_lwes, streams.stream(0),
streams.gpu_index(0),
gpu_memory_allocated);
cuda_drop_with_size_tracking_async(lwe_trivial_indexes, streams.stream(0),
streams.gpu_index(0),
gpu_memory_allocated);
cuda_drop_with_size_tracking_async(d_expand_jobs, streams.stream(0),
streams.gpu_index(0),
gpu_memory_allocated);