From d26498627cc0d52ba5dd3600e148432b7b70f8ff Mon Sep 17 00:00:00 2001
From: willzhang4a58 <willzhang4a58@gmail.com>
Date: Tue, 13 Jun 2017 15:55:08 +0800
Subject: [PATCH] total_reading_cnt_

---
 oneflow/core/actor/actor.cpp | 3 +++
 oneflow/core/actor/actor.h   | 2 ++
 2 files changed, 5 insertions(+)

diff --git a/oneflow/core/actor/actor.cpp b/oneflow/core/actor/actor.cpp
index 9564b73a21..376c040a06 100644
--- a/oneflow/core/actor/actor.cpp
+++ b/oneflow/core/actor/actor.cpp
@@ -41,6 +41,7 @@ void Actor::Init(const TaskProto& task_proto) {
     produced_regst2reading_cnt_[regst.get()] = 0;
   }
   writeable_produced_regst_desc_num_ = writeable_produced_regst_.size();
+  total_reading_cnt_ = 0;
 }
 
 void Actor::AsyncWardKernelAndSendMsgToRegstReader(
@@ -64,6 +65,7 @@ void Actor::AsyncWardKernelAndSendMsgToRegstReader(
       }
     });
     produced_regst2reading_cnt_.at(regst) = regst->subscribers_actor_id().size();
+    total_reading_cnt_ += regst->subscribers_actor_id().size();
     if (!regst->subscribers_actor_id().empty()) { pair.second.pop(); }
     if (pair.second.empty()) { writeable_produced_regst_desc_num_ -= 1; }
   }
@@ -74,6 +76,7 @@ int Actor::TryUpdtStateAsFromRegstReader(Regst* regst) {
   if (reading_cnt_it == produced_regst2reading_cnt_.end()) { return -1; }
   CHECK_GE(reading_cnt_it->second, 1);
   reading_cnt_it->second -= 1;
+  total_reading_cnt_ -= 1;
   if (reading_cnt_it->second != 0) { return 0; }
   auto writeable_it = writeable_produced_regst_.find(regst->regst_desc_id());
   if (writeable_it == writeable_produced_regst_.end()) { return 0; }
diff --git a/oneflow/core/actor/actor.h b/oneflow/core/actor/actor.h
index 3cdd55ddf1..2814002a94 100644
--- a/oneflow/core/actor/actor.h
+++ b/oneflow/core/actor/actor.h
@@ -54,6 +54,7 @@ class Actor {
     if (!it->second.empty()) { writeable_produced_regst_desc_num_ -= 1; }
     writeable_produced_regst_.erase(it);
   }
+  int64_t total_reading_cnt() const { return total_reading_cnt_; }
 
  private:
   uint64_t actor_id_;
@@ -67,6 +68,7 @@ class Actor {
   HashMap<uint64_t, std::queue<Regst*>> writeable_produced_regst_; // <regst_desc_id, regst>
   uint64_t writeable_produced_regst_desc_num_;
   HashMap<Regst*, int64_t> produced_regst2reading_cnt_;
+  int64_t total_reading_cnt_;
 
 };
 
-- 
GitLab