Pre Merge pull request !2080 from wuweikang/device_id

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;
     }

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;
         }
       }

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);

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);