From d2cf7c0de52d45b53dc26df0aef6d61f298ac5fd Mon Sep 17 00:00:00 2001
From: wuweikang <wuweikang@huawei.com>
Date: Tue, 31 Aug 2021 10:21:06 +0800
Subject: [PATCH] add device id when initialize and finalize

---
 ge/graph/manager/graph_mem_manager.cc   | 2 +-
 ge/graph/manager/graph_mem_manager.h    | 2 +-
 ge/graph/manager/rdma_pool_allocator.cc | 7 ++++---
 ge/graph/manager/rdma_pool_allocator.h  | 5 +++--
 4 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/ge/graph/manager/graph_mem_manager.cc b/ge/graph/manager/graph_mem_manager.cc
index 21eaf302..8030c970 100644
--- a/ge/graph/manager/graph_mem_manager.cc
+++ b/ge/graph/manager/graph_mem_manager.cc
@@ -73,7 +73,7 @@ template <typename T>
 void FinalizeAllocatorMap(std::map<rtMemType_t, T *> &allocate_map) {
   for (auto &allocator : allocate_map) {
     if (allocator.second != nullptr) {
-      allocator.second->Finalize();
+      allocator.second->Finalize(GetContext().DeviceId());
       delete allocator.second;
       allocator.second = nullptr;
     }
diff --git a/ge/graph/manager/graph_mem_manager.h b/ge/graph/manager/graph_mem_manager.h
index d7993ed4..2d142d83 100644
--- a/ge/graph/manager/graph_mem_manager.h
+++ b/ge/graph/manager/graph_mem_manager.h
@@ -96,7 +96,7 @@ class MemManager {
         GELOGE(ACL_ERROR_GE_MEMORY_ALLOCATION, "Create Allocator failed.");
         return ACL_ERROR_GE_MEMORY_ALLOCATION;
       } else {
-        if (allocator->Initialize() != SUCCESS) {
+        if (allocator->Initialize(GetContext().DeviceId()) != SUCCESS) {
           return ACL_ERROR_GE_INTERNAL_ERROR;
         }
       }
diff --git a/ge/graph/manager/rdma_pool_allocator.cc b/ge/graph/manager/rdma_pool_allocator.cc
index 4297be95..7e4b46a1 100644
--- a/ge/graph/manager/rdma_pool_allocator.cc
+++ b/ge/graph/manager/rdma_pool_allocator.cc
@@ -49,15 +49,16 @@ RdmaPoolAllocator::RdmaPoolAllocator(rtMemType_t memory_type)
         return reinterpret_cast<uintptr_t>(left->ptr) < reinterpret_cast<uintptr_t>(right->ptr);
       })) {}
 
-Status RdmaPoolAllocator::Initialize() {
+Status RdmaPoolAllocator::Initialize(uint32_t device_id) {
+  GELOGI("Device id %u", device_id);
   memory_allocator_ = &MemManager::Instance().MemInstance(memory_type_);
   if (memory_allocator_ == nullptr) {
     return ACL_ERROR_GE_INTERNAL_ERROR;
   }
   return ge::SUCCESS;
 }
-void RdmaPoolAllocator::Finalize() {
-  GELOGD("Rdma pool finalize start.");
+void RdmaPoolAllocator::Finalize(uint32_t device_id) {
+  GELOGD("Rdma pool finalize start, device id:%u", device_id);
   for (auto it = allocated_blocks_.begin(); it != allocated_blocks_.end();) {
     auto block = it->second;
     it = allocated_blocks_.erase(it);
diff --git a/ge/graph/manager/rdma_pool_allocator.h b/ge/graph/manager/rdma_pool_allocator.h
index 0a895a11..5f06e9ba 100644
--- a/ge/graph/manager/rdma_pool_allocator.h
+++ b/ge/graph/manager/rdma_pool_allocator.h
@@ -42,8 +42,9 @@ class RdmaPoolAllocator {
 
   ~RdmaPoolAllocator() = default;
 
-  Status Initialize();
-  void Finalize();
+  Status Initialize(uint32_t device_id = 0);
+
+  void Finalize(uint32_t device_id = 0);
 
   Status InitMemory(size_t mem_size);