From 4f9f4982f6804f5f0db7cb7bacafcc72297c9f98 Mon Sep 17 00:00:00 2001 From: Agnes Leroy Date: Fri, 14 Nov 2025 11:07:27 +0100 Subject: [PATCH] fix(gpu): fix memory leak in rerand --- .../cuda/include/integer/rerand_utilities.h | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/backends/tfhe-cuda-backend/cuda/include/integer/rerand_utilities.h b/backends/tfhe-cuda-backend/cuda/include/integer/rerand_utilities.h index f72c1d91f..6d3c908e0 100644 --- a/backends/tfhe-cuda-backend/cuda/include/integer/rerand_utilities.h +++ b/backends/tfhe-cuda-backend/cuda/include/integer/rerand_utilities.h @@ -7,10 +7,10 @@ template 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 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 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);