hashtable save and load

Author: houzhenggang

fbgemm_gpu/src/dram_kv_embedding_cache/SynchronizedShardedMap.h

@@ -75,9 +75,10 @@ class SynchronizedShardedMap {
     out.close();
 
     // save every mempool
-    for (std::size_t i = 0; i < num_shards; ++i) {
+    for (std::size_t shard_id = 0; shard_id < getNumShards(); ++shard_id) {
       std::string pool_filename = filename + ".pool." + std::to_string(i);
-      mempools_[i]->serialize(pool_filename);
+      auto wlock = shards_[shard_id].wlock();
+      mempools_[shard_id]->serialize(pool_filename);
     }
   }
 
@@ -95,18 +96,14 @@ class SynchronizedShardedMap {
       throw std::runtime_error("Shard count mismatch between file and map");
     }
 
-    // first deserialize mempool
-    for (std::size_t i = 0; i < num_shards; ++i) {
-      std::string pool_filename = filename + ".pool." + std::to_string(i);
-      mempools_[i]->deserialize(pool_filename);
-    }
-
-    // load map from mempool
     for (std::size_t shard_id = 0; shard_id < getNumShards(); ++shard_id) {
+      std::string pool_filename = filename + ".pool." + std::to_string(i);
       auto wlock = shards_[shard_id].wlock();
-      auto* mempool = mempools_[shard_id].get();
+      // first deserialize mempool
+      mempools_[shard_id]->deserialize(pool_filename);
+      // load map from mempool
       wlock->clear();
-      mempool->for_each_block([&wlock](void* block) {
+      mempools_[shard_id]->for_each_block([&wlock](void* block) {
         auto key = FixedBlockPool::get_key(block);
         wlock->emplace(key, reinterpret_cast<V>(block));
       });