Skip to content
Snippets Groups Projects
Commit d26a6fce authored by Li Xinqi's avatar Li Xinqi Committed by Jinhui Yuan
Browse files

Dev used order value when mem sharing (#940)

* is head regst desc where sharing same mem

* fix typo

* memory sharing info
parent 16ce3282
No related branches found
No related tags found
No related merge requests found
......@@ -58,6 +58,7 @@ void CopyCommNetTaskNode::Init(int64_t machine_id, int64_t peer_machine_id) {
}
namespace {
HashMap<int64_t, HashMap<int64_t, int64_t>>* GetConnection2LocalStreamIdMap() {
// this_machine_id -> {peer_machine_id, local_work_stream_id}
static HashMap<int64_t, HashMap<int64_t, int64_t>> connection2stream_id;
......@@ -77,6 +78,7 @@ void InsertLocalStreamId4Connection(int64_t this_machine_id, int64_t peer_machin
auto& dict = *GetConnection2LocalStreamIdMap();
dict[this_machine_id][peer_machine_id] = dict[this_machine_id].size();
}
} // namespace
int64_t CopyCommNetTaskNode::AllocateLocalWorkStreamId() {
......
......@@ -22,7 +22,8 @@ RegstDesc::RegstDesc() {
max_register_num_ = kMaxRegisterNum;
is_locked_ = false;
enable_mem_sharing_ = false;
mem_shared_id_ = -1;
mem_sharing_info_.set_mem_shared_id(-1);
mem_sharing_info_.set_used_order_value(-1);
}
void RegstDesc::AddConsumer(const TaskNode* new_consumer) {
......@@ -132,9 +133,7 @@ void RegstDesc::ToProto(RegstDescProto* ret) const {
ret->set_max_register_num(max_register_num_);
ret->set_register_num(min_register_num_);
*(ret->mutable_mem_case()) = mem_case_;
ret->set_enable_mem_sharing(enable_mem_sharing_);
ret->set_mem_shared_id(mem_shared_id_);
ret->set_is_head_when_sharing_mem(false);
if (enable_mem_sharing_) { *ret->mutable_mem_sharing_info() = mem_sharing_info_; }
}
bool RegstDesc::HasSameBlobDescs(const RegstDesc* rhs) {
......
......@@ -68,7 +68,7 @@ class RegstDesc final {
MemoryCase mem_case_;
RegstDescTypeProto regst_desc_type_;
bool enable_mem_sharing_;
int32_t mem_shared_id_;
MemSharingProto mem_sharing_info_;
};
} // namespace oneflow
......
......@@ -31,6 +31,11 @@ message RegstDescTypeProto {
}
}
message MemSharingProto {
required int32 mem_shared_id = 1;
required int32 used_order_value = 2;
}
message RegstDescProto {
required int64 regst_desc_id = 1;
required int64 producer_task_id = 2;
......@@ -40,7 +45,5 @@ message RegstDescProto {
required int32 register_num = 6;
required MemoryCase mem_case = 7;
required RegstDescTypeProto regst_desc_type = 8;
required bool enable_mem_sharing = 9;
required int32 mem_shared_id = 10;
required bool is_head_when_sharing_mem = 11;
optional MemSharingProto mem_sharing_info = 9;
}
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment