From 0329d4fc2dfa72ba2747ef24cf874695b43aafd9 Mon Sep 17 00:00:00 2001 From: Antoniu Pop Date: Tue, 17 Jan 2023 22:29:12 +0000 Subject: [PATCH] fix(gpu-wrappers): fix KS/BS wrappers for GPU (memory management). --- compiler/lib/Runtime/wrappers.cpp | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/compiler/lib/Runtime/wrappers.cpp b/compiler/lib/Runtime/wrappers.cpp index 904d96f3c..4318241e8 100644 --- a/compiler/lib/Runtime/wrappers.cpp +++ b/compiler/lib/Runtime/wrappers.cpp @@ -149,7 +149,6 @@ void memref_batched_keyswitch_lwe_cuda_u64( // free memory that we allocated on gpu cuda_drop(ct0_gpu, gpu_idx); cuda_drop(out_gpu, gpu_idx); - cuda_drop(ksk_gpu, gpu_idx); cuda_destroy_stream(stream, gpu_idx); } @@ -197,10 +196,7 @@ void memref_batched_bootstrap_lwe_cuda_u64( // Move the glwe accumulator to the GPU void *glwe_ct_gpu = - alloc_and_memcpy_async_to_gpu(glwe_ct, 0, glwe_ct_size, gpu_idx, stream); - - // Free the glwe accumulator (on CPU) - free(glwe_ct); + alloc_and_memcpy_async_to_gpu(glwe_ct, 0, glwe_ct_len, gpu_idx, stream); // Move test vector indexes to the GPU, the test vector indexes is set of 0 uint32_t num_test_vectors = 1, lwe_idx = 0, @@ -220,6 +216,8 @@ void memref_batched_bootstrap_lwe_cuda_u64( memcpy_async_to_cpu(out_aligned, out_offset, out_batch_size, out_gpu, gpu_idx, stream); cuda_synchronize_device(gpu_idx); + // Free the glwe accumulator (on CPU) + free(glwe_ct); // free memory that we allocated on gpu cuda_drop(ct0_gpu, gpu_idx); cuda_drop(out_gpu, gpu_idx);