expected piece id

oneflow/core/actor/actor.cpp

@@ -33,6 +33,8 @@ void Actor::Init(const TaskProto& task_proto) {
   for (const auto& pair : task_proto.subscribed_regst_desc_id()) {
     CHECK(name2regst_desc_id_.emplace(pair.first, pair.second).second);
   }
+  //
+  expected_piece_id_ = 0;
   // Status of Produced Registers
   for (const auto& regst : produced_regst_vec_) {
     writeable_produced_regst_[regst->regst_desc_id()].push(regst.get());
@@ -52,6 +54,7 @@ void Actor::WardKernel(
       return regst->GetBlobPtrFromLbn(lbn);
     });
   }
+  expected_piece_id_ += 1;
 }
 
 void Actor::ForEachProducedRegst(std::function<void(Regst*)> func) {

oneflow/core/actor/actor.h

@@ -39,6 +39,7 @@ class Actor {
     return name2regst_desc_id_.at(name);
   }
 
+  uint64_t expected_piece_id() const { return expected_piece_id_; }
   // Status of Produced Registers
   int TryOneReadDone(Regst* regst);
   Regst* GetCurWriteableRegst(uint64_t regst_desc_id);
@@ -59,6 +60,7 @@ class Actor {
   std::vector<std::unique_ptr<Regst>> produced_regst_vec_;
   HashMap<std::string, uint64_t> name2regst_desc_id_;
   
+  uint64_t expected_piece_id_;
   // Status of Produced Registers
   HashMap<uint64_t, std::queue<Regst*>> writeable_produced_regst_; // <regst_desc_id, regst>
   uint64_t writeable_produced_regst_desc_num_;