improve readability

oneflow/register/register_manager.cpp

@@ -6,9 +6,9 @@
 namespace oneflow {
 
 void RegstMgr::NewRegsts(const RegstDescProto& regst_desc_proto,
-               std::function<void(Regst*)> OneRegstDone) {
+                         std::function<void(Regst*)> OneRegstDone) {
   std::shared_ptr<RtRegstDesc> runtime_regst_desc(new RtRegstDesc(regst_desc_proto));
-  for (size_t i = 0; i < regst_desc_proto.register_num(); ++i) {
+  for (uint64_t i = 0; i < regst_desc_proto.register_num(); ++i) {
     Regst* regst = new Regst;
     regst->regst_desc_ = runtime_regst_desc;
     regst->regst_id_ =
@@ -19,9 +19,10 @@ void RegstMgr::NewRegsts(const RegstDescProto& regst_desc_proto,
       elem_size = sizeof(double);
     }
     uint64_t elem_cnt = 0;
-    std::vector<std::string> lbns(regst_desc_proto.lbn2shape().size());
+    std::vector<std::string> lbns;
+    lbns.reserve(regst_desc_proto.lbn2shape().size());
     for (const auto& pair : regst_desc_proto.lbn2shape()) {
-      Shape* shape_ptr = runtime_regst_desc->GetShapePtrFromLbn(pair.first);
+      const Shape* shape_ptr = runtime_regst_desc->GetShapePtrFromLbn(pair.first);
       lbns.push_back(pair.first);
       elem_cnt += shape_ptr->elem_cnt();
     }
@@ -32,8 +33,9 @@ void RegstMgr::NewRegsts(const RegstDescProto& regst_desc_proto,
 
     uint64_t blob_idx = 0;
     for (const std::string& lbn : lbns) {
-      Shape* shape_ptr = runtime_regst_desc->GetShapePtrFromLbn(lbn);
-      CHECK(regst->lbn2blob_.emplace(lbn, of_make_unique<Blob>(allocation.first + blob_idx, shape_ptr)).second);
+      const Shape* shape_ptr = runtime_regst_desc->GetShapePtrFromLbn(lbn);
+      auto blob_ptr = of_make_unique<Blob>(allocation.first + blob_idx, shape_ptr);
+      CHECK(regst->lbn2blob_.emplace(lbn, std::move(blob_ptr)).second);
       blob_idx += shape_ptr->elem_cnt() * elem_size;
     }
     regst->deleter_ = allocation.second;

oneflow/register/runtime_register_desc.cpp

+#include "register/runtime_register_desc.h"
+#include "common/id_manager.h"
+
+namespace oneflow {
+
+RtRegstDesc::RtRegstDesc(const RegstDescProto& regst_desc_proto) {
+  regst_desc_id_ = regst_desc_proto.regst_desc_id();
+  producer_actor_id_ = 
+    IDMgr::Singleton().GetActorIdFromTaskId(regst_desc_proto.producer_task_id());
+  register_num_ = regst_desc_proto.register_num();
+
+  const auto& subscriber = regst_desc_proto.subscriber_task_id();
+  subscribers_actor_id_.reserve(subscriber.size());
+  for (uint64_t task_id : subscriber) {
+    subscribers_actor_id_.push_back(IDMgr::Singleton().GetActorIdFromTaskId(task_id));
+  }
+
+  for (const auto& pair : regst_desc_proto.lbn2shape()) {
+    lbn2shape_.emplace(pair.first, of_make_unique<Shape>(pair.second));
+  }
+  mem_case_ = regst_desc_proto.mem_case();
+}
+
+} // namespace oneflow

oneflow/register/runtime_register_desc.h

@@ -2,9 +2,7 @@
 #define ONEFLOW_REGISTER_RUNTIME_REGISTER_DESC_H_
 
 #include "common/util.h"
-#include "common/protobuf.h"
 #include "common/shape.h"
-#include "common/id_manager.h"
 #include "memory/memory_case.pb.h"
 #include "register/register_desc.pb.h"
 
@@ -16,23 +14,7 @@ class RtRegstDesc {
   RtRegstDesc() = delete;
   ~RtRegstDesc() = default;
 
-  RtRegstDesc(const RegstDescProto& regst_desc_proto) {
-    regst_desc_id_ = regst_desc_proto.regst_desc_id();
-    producer_actor_id_ = 
-      IDMgr::Singleton().GetActorIdFromTaskId(regst_desc_proto.producer_task_id());
-    register_num_ = regst_desc_proto.register_num();
-
-    const auto& subscriber = regst_desc_proto.subscriber_task_id();
-    subscribers_actor_id_.reserve(subscriber.size());
-    for (uint64_t task_id : subscriber) {
-      subscribers_actor_id_.push_back(IDMgr::Singleton().GetActorIdFromTaskId(task_id));
-    }
-
-    for (const auto& pair : regst_desc_proto.lbn2shape()) {
-      lbn2shape_.emplace(pair.first, of_make_unique<Shape>(pair.second));
-    }
-    mem_case_ = regst_desc_proto.mem_case();
-  }
+  RtRegstDesc(const RegstDescProto& regst_desc_proto);
 
   uint64_t regst_desc_id() const { return regst_desc_id_; }
   uint64_t producer_actor_id() const { return producer_actor_id_; }
@@ -42,7 +24,7 @@ class RtRegstDesc {
   int64_t register_num() const { return register_num_; }
   const MemoryCase& mem_case() const { return mem_case_; }
 
-  Shape* GetShapePtrFromLbn(const std::string& lbn) {
+  const Shape* GetShapePtrFromLbn(const std::string& lbn) const {
     return lbn2shape_.at(lbn).get();
   }