How to free memory allocated by SEAL?

[crockeea]

For testing and benchmarking, I wrote a loop which allocates several different CKKS parameters and generates keys for those parameters. I expected that all memory allocated for the keys/params is freed at the end of the for loop. Instead, this memory is apparently not freed.

#include "seal/seal.h"

using namespace std;
using namespace seal;

vector<int> gen_modulus_vec(int num_primes, int log_scale) {
  vector<int> modulusVector(num_primes);
  modulusVector[0] = 60;
  for (int i = 1; i < num_primes - 1; i++) {
    modulusVector[i] = log_scale;
  }
  modulusVector[num_primes - 1] = 60;
  return modulusVector;
}

int main() {
  double log_scale = 40;
  int numSlots = 16384;

  for(int i = 1; i < 7; i++) {
    EncryptionParameters parms(scheme_type::CKKS);
    size_t poly_modulus_degree = 2*numSlots;
    parms.set_poly_modulus_degree(poly_modulus_degree);
    parms.set_coeff_modulus(CoeffModulus::Create(poly_modulus_degree, gen_modulus_vec(3*i+2, log_scale)));
    auto context = SEALContext::Create(parms);

    KeyGenerator keygen(context);
    auto public_key = keygen.public_key();
    auto secret_key = keygen.secret_key();
    auto relin_keys = keygen.relin_keys_local();
    auto gal_keys = keygen.galois_keys_local();
  }
}

I tracked memory usage at the start of each for iteration:
Start of i=1: 905M
Start of i=2: 1.2G
Start of i=3: 2.1G
Start of 1=4: 3.7G
Start of i=5: 6.4G
Start of i=6: 10G

How can I free memory allocated by SEAL?

[kimlaine]

When you create new parameter sets, the buffer allocation sizes change totally. SEAL memory pool only reuses memory allocations if their sizes match exactly. The justification for that is that in many cases we would expect an application reuse the same set of parameters multiple times, and with exact size match we can ensure there is no "wasted space". However, now you run into this problem.

All allocations that you are making are taken from the so-called global memory pool. For example, the constructor of GaloisKeys (inheriting from KSwitchKeys) will acquire a memory pool for its internal storage with MemoryManager::GetPool(), which acquires the pool according to the currently set memory manager profile. The default profile always allocates from the global memory pool, so what you need to do is change the profile so the pool is some custom one. If you don't hold any extra references to your custom memory pool, it will be released automatically once there are no more references pointing to it.

Since your code doesn't necessarily make sense as a real application, it's a little hard to say how exactly you should do this, but one way is as follows. Change the beginning of your loop to:

for(int i = 1; i < 7; i++) {
   MemoryPoolHandle my_pool = MemoryPoolHandle::New();
   MemoryManager::SwitchProfile(make_unique<MMProfFixed>(std::move(my_pool)));

   EncryptionParameters parms(scheme_type::CKKS);
   // Your code here ...

Here you first create a new memory pool. Then you switch the memory manager profile to a "fixed" profile that always returns the same pool whenever MemoryManager::GetPool() is called. In this case it always returns your my_pool, which then is used in all allocations for that iteration of the loop.

Once the loop completes, however, you may want to switch back to the previous profile (the one that always returns the global memory pool). This can be done as follows:

for(int i = 1; i < 7; i++) {
   MemoryPoolHandle my_pool = MemoryPoolHandle::New();
   auto old_prof = MemoryManager::SwitchProfile(make_unique<MMProfFixed>(std::move(my_pool)));

   EncryptionParameters parms(scheme_type::CKKS);
   // Your code here ...

   auto gal_keys = keygen.galois_keys_local();
   MemoryManager::SwitchProfile(std::move(old_prof));
}

You can also use the MMProfGuard class to achieve the same effect more easily:

for(int i = 1; i < 7; i++) {
   MemoryPoolHandle my_pool = MemoryPoolHandle::New();
   auto pg = MemoryManager::MMProfGuard(make_unique<MMProfFixed>(std::move(my_pool)));

   EncryptionParameters parms(scheme_type::CKKS);
   // Your code here ...

   auto gal_keys = keygen.galois_keys_local();
}

If you need some custom memory pool behavior, it's easy to create your own memory manager profile (just inherit from MMProf). You can also create your custom memory pool if you want total control. The MemoryManager::GetPool(...) function takes a mm_prof_opt_t (just a std::uint64_t) that it passes to the profile; a custom profile can then optionally decide to use it for some internal logic to decide what kind of pool to return.

[crockeea]

Kim, thanks for your detailed answer.
Your second snippet produces the memory usage profile that I expected: the usage at the beginning of each iteration of the loop always the same small amount, while the usage at the end of the loop gets larger with each iteration (due to increasing parameter sizes).

However, I'd appreciate your help understanding what's happening with the 1st and 3rd snippets in your answer. The first snippet is the same as the second, except that you don't reset to use the global memory pool at the end of the loop. When I do that, I get the an odd memory usage profile that is different from the original, but which still leaks memory over time. Specifically:

for(int i = 1; i < 7; i++) {
    // measure mem usage here

    MemoryPoolHandle my_pool = MemoryPoolHandle::New();
    MemoryManager::SwitchProfile(make_unique<MMProfFixed>(std::move(my_pool)));

    EncryptionParameters parms(scheme_type::CKKS);
    size_t poly_modulus_degree = 2*numSlots;
    parms.set_poly_modulus_degree(poly_modulus_degree);
    parms.set_coeff_modulus(CoeffModulus::Create(poly_modulus_degree, gen_modulus_vec(3*i+2, log_scale)));
    auto context = SEALContext::Create(parms);

    KeyGenerator keygen(context);
    auto public_key = keygen.public_key();
    auto secret_key = keygen.secret_key();
    auto relin_keys = keygen.relin_keys_local();
    auto gal_keys = keygen.galois_keys_local();

    // measure mem usage here
  }

Memory usage at beginning and end of each loop iteration:

Loop Index               Mem use at beginning of loop              Mem us at end of Loop
1                                 859M                                      1.2G
2                                 1.2G                                      2.1G
3                                 1.7G                                      3.5G
4                                 2.6G                                      5.4G
5                                 3.9G                                      7.9G
6                                 5.4G                                      10G

So you can see that some memory is being freed at the end of each loop (when my_pool is released), but some of the memory is retained.

For your 3rd snippet, I'm assuming you meant to write auto pg = MMProfGuard(make_unique<MMProfFixed>(std::move(my_pool)));. Since you didn't use pg anywhere, I just used MMProfGuard(make_unique<MMProfFixed>(std::move(my_pool)));, and still got the memory leak. Using your suggested auto pg = MMProfGuard(make_unique<MMProfFixed>(std::move(my_pool)));, I get the "good" memory usage profile, so it seems like MMProfGuard depends on an object being created and destroyed at the end of a scope. I was surprised that it didn't generate an unused variable compiler warning; there must be some magic going on.

[WeiDaiWD]

Sorry that this issue has been here for such a long time. I've tried to reproduce this issue with version 3.7.1. It seems like that it has been resolved overtime. At least from my side I don't see a memory leakage. Do you still have this issue?