Add blocking counter

paddle/framework/CMakeLists.txt

@@ -60,6 +60,7 @@ cc_library(selected_rows SRCS selected_rows.cc DEPS tensor)
 cc_test(selected_rows_test SRCS selected_rows_test.cc DEPS selected_rows)
 
 cc_test(threadpool_test SRCS threadpool_test.cc)
+cc_test(blocking_counter_test SRCS blocking_counter_test.cc)
 cc_library(init SRCS init.cc DEPS gflags device_context place stringpiece)
 cc_test(init_test SRCS init_test.cc DEPS init)
 

paddle/framework/blocking_counter.h

@@ -0,0 +1,68 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+#include <atomic>
+#include <condition_variable>
+#include <cstdio>
+#include <functional>
+#include <mutex>
+#include <thread>
+
+#include "paddle/platform/call_once.h"
+#include "paddle/platform/enforce.h"
+
+namespace paddle {
+namespace framework {
+
+class BlockingCounter {
+ public:
+  BlockingCounter(int cnt) : cnt_(cnt), done_(false) {
+    PADDLE_ENFORCE_GE(static_cast<int>(cnt_), 0,
+                      "The initialized counter should not be less than 0");
+  }
+
+  ~BlockingCounter() {}
+
+  /**
+   * @breif   Wait untile the count was decreased to 0
+   */
+  void Wait() {
+    if (cnt_ == 0) return;
+    std::unique_lock<std::mutex> lock(done_m);
+    done_cv.wait(lock, [=] { return done_ == true; });
+  }
+
+  /**
+   * @breif   Descrease the count by 1
+   */
+  void DecreaseCount() {
+    cnt_.fetch_sub(1);
+    if (cnt_ != 0) return;
+    std::unique_lock<std::mutex> lock(done_m);
+    done_ = true;
+    lock.unlock();
+    done_cv.notify_all();
+  }
+
+ private:
+  BlockingCounter(const BlockingCounter&) = delete;
+  std::atomic<int> cnt_;
+  bool done_;
+  std::mutex done_m;
+  std::condition_variable done_cv;
+};
+
+}  // namespace framework
+}  // namespace paddle

paddle/framework/blocking_counter_test.cc

@@ -0,0 +1,54 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "blocking_counter.h"
+#include <gtest/gtest.h>
+#include <atomic>
+#include <thread>
+#include "threadpool.h"
+
+namespace framework = paddle::framework;
+
+TEST(BlockingCounter, SingleThread) {
+  framework::BlockingCounter bc(2);
+  bc.DecreaseCount();
+  bc.DecreaseCount();
+  bc.Wait();
+}
+
+TEST(BlockingCounter, MultipleThread) {
+  framework::ThreadPool* pool = framework::ThreadPool::GetInstance();
+  std::vector<int> cnt_per_task = {10, 2000, 30, 40, 50};
+  int task_cnt = cnt_per_task.size();
+  std::vector<framework::BlockingCounter*> bcs(task_cnt);
+  std::vector<std::atomic<int>> sums(task_cnt);
+  for (int i = 0; i < task_cnt; ++i) {
+    int cnt = cnt_per_task[i];
+    auto& sum = sums[i];
+    sum = 0;
+    auto& bc = bcs[i];
+    bc = new framework::BlockingCounter(cnt);
+    pool->Run([&sum, bc, cnt]() {
+      for (int j = 0; j < cnt; ++j) {
+        sum.fetch_add(1);
+        bc->DecreaseCount();
+      }
+    });
+    bcs.push_back(std::move(bc));
+  }
+
+  for (int i = 0; i < task_cnt; ++i) {
+    bcs[i]->Wait();
+  }
+}