RegstDesc4RegstDescId (#1277)

Former-commit-id: b1987003

oneflow/core/register/register_manager.cpp

@@ -155,4 +155,10 @@ void RegstMgr::InitOFRecordBlobIfNeed(Blob* blob_ptr) {
   }
 }
 
+const RtRegstDesc& RegstMgr::RegstDesc4RegstDescId(int64_t regst_desc_id) const {
+  const auto& it = regst_desc_id2rt_regst_desc_.find(regst_desc_id);
+  CHECK(it != regst_desc_id2rt_regst_desc_.end());
+  return *it->second;
+}
+
 }  // namespace oneflow

oneflow/core/register/register_manager.h

@@ -18,6 +18,7 @@ class RegstMgr final {
   ~RegstMgr() = default;
 
   void NewRegsts(const RegstDescProto& regst_desc_proto, std::function<void(Regst*)> OneRegstDone);
+  const RtRegstDesc& RegstDesc4RegstDescId(int64_t regst_desc_id) const;
 
  private:
   friend class Global<RegstMgr>;