Merge branch 'gpugraph_new' of https://github.com/xuewujiao/Paddle in…

…to fix_600_infer
xuewujiao · May 16, 2023 · fdaefb2 · fdaefb2
2 parents 58671ba + f100887
commit fdaefb2
Show file tree

Hide file tree

Showing 7 changed files with 602 additions and 607 deletions.
diff --git a/paddle/fluid/distributed/ps/service/ps_local_client.cc b/paddle/fluid/distributed/ps/service/ps_local_client.cc
@@ -33,6 +33,9 @@ int32_t PsLocalClient::Initialize() {
 
 ::std::future<int32_t> PsLocalClient::Shrink(uint32_t table_id,
                                              const std::string threshold) {
+  // threshold not use
+  auto* table_ptr = GetTable(table_id);
+  table_ptr->Shrink("");
   return done();
 }
 

diff --git a/paddle/fluid/distributed/ps/table/memory_sparse_table.cc b/paddle/fluid/distributed/ps/table/memory_sparse_table.cc
@@ -1050,18 +1050,26 @@ int32_t MemorySparseTable::Flush() { return 0; }
 
 int32_t MemorySparseTable::Shrink(const std::string &param) {
   VLOG(0) << "MemorySparseTable::Shrink";
-  // TODO(zhaocaibei123): implement with multi-thread
+  std::atomic<uint32_t> shrink_size_all{0};
+  int thread_num = _real_local_shard_num;
+  omp_set_num_threads(thread_num);
+#pragma omp parallel for schedule(dynamic)
   for (int shard_id = 0; shard_id < _real_local_shard_num; ++shard_id) {
     // Shrink
+    int feasign_size = 0;
     auto &shard = _local_shards[shard_id];
     for (auto it = shard.begin(); it != shard.end();) {
       if (_value_accesor->Shrink(it.value().data())) {
         it = shard.erase(it);
+        ++feasign_size;
       } else {
         ++it;
       }
     }
+    shrink_size_all += feasign_size;
   }
+  VLOG(0) << "MemorySparseTable::Shrink success, shrink size:"
+          << shrink_size_all;
   return 0;
 }