Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #include "log.h"
      12              : #include "workspace_mem.h"
      13              : 
      14              : namespace hccl {
      15              : 
      16          524 : WorkSpaceMem::WorkSpaceMem() = default;
      17              : 
      18          524 : WorkSpaceMem::~WorkSpaceMem() = default;
      19              : 
      20              : /* 适配GE，设备内存由上层分配，HCCL主要用于管理该设备内存 */
      21          250 : HcclResult WorkSpaceMem::SetMemResource(const std::string& tag, void* ptr, u64 maxSize)
      22              : {
      23          250 :     if (ptr == nullptr) {
      24            0 :         HCCL_ERROR("[Set][MemResource]construct fail, ptr is [null], tag[%s], maxSize[%llu]", tag.c_str(), maxSize);
      25            0 :         return HCCL_E_PTR;
      26              :     } else {
      27          250 :         std::unique_lock<std::mutex> lock(memResMutex_);
      28          250 :         WorkSpaceMem_t memResource{};
      29          250 :         memResource.curPtr = ptr;
      30          250 :         memResource.maxSize = maxSize;
      31          250 :         memResMap_.erase(tag);
      32          250 :         memResMap_.insert(std::make_pair(tag, memResource));
      33          250 :         lock.unlock();
      34          250 :         HCCL_INFO(
      35              :             "[Set][MemResource]set mem resource success, tag[%s] ptr[%llu] max size[%llu]", tag.c_str(),
      36              :             std::hash<void*>{}(ptr), maxSize);
      37          250 :     }
      38          250 :     return HCCL_SUCCESS;
      39              : }
      40              : 
      41           23 : void* WorkSpaceMem::AllocMem(const std::string& tag, u64 size)
      42              : {
      43              :     /* 此处可能会与并发，加锁 */
      44           23 :     std::unique_lock<std::mutex> lock(memResMutex_);
      45              : 
      46           24 :     auto interIter = memResMap_.find(tag);
      47           24 :     if (interIter == memResMap_.end()) {
      48           12 :         HCCL_ERROR("[Alloc][Mem] Can't find tag, tag[%s], size[%llu]", tag.c_str(), size);
      49           12 :         return nullptr;
      50              :     }
      51           12 :     if (interIter->second.curPtr == nullptr) {
      52            0 :         HCCL_ERROR("[Alloc][Mem] Current resource(curPtr) is null, tag[%s], size[%llu]", tag.c_str(), size);
      53            0 :         return nullptr;
      54              :     }
      55              : 
      56              :     // 判断size 是否超出最大值
      57           12 :     if (interIter->second.maxSize < (interIter->second.totalSize + size)) {
      58            0 :         HCCL_ERROR(
      59              :             "[Alloc][Mem] Current resource size is not enough, tag[%s], size[%llu], "
      60              :             "maxSize[%llu], totalSize[%llu]",
      61              :             tag.c_str(), size, interIter->second.maxSize, interIter->second.totalSize);
      62            0 :         return nullptr;
      63              :     }
      64              : 
      65           11 :     void* tempPtr = interIter->second.curPtr;
      66            9 :     interIter->second.curPtr = reinterpret_cast<void*>(static_cast<char*>(interIter->second.curPtr) + size);
      67            9 :     interIter->second.totalSize += size;
      68              : 
      69            9 :     HCCL_INFO(
      70              :         "[Alloc][Mem]Alloc mem success, tag[%s] size[%llu], totalSize[%llu], maxSize[%llu]", tag.c_str(), size,
      71              :         interIter->second.totalSize, interIter->second.maxSize);
      72           12 :     lock.unlock();
      73           12 :     return tempPtr;
      74           24 : }
      75              : 
      76           15 : HcclResult WorkSpaceMem::DestroyMemResource(const std::string& tag)
      77              : {
      78           15 :     HCCL_INFO("[Destroy][MemResource]Destroy workspace mem, tag[%s]", tag.c_str());
      79              : 
      80           16 :     std::unique_lock<std::mutex> lock(memResMutex_);
      81              : 
      82           14 :     auto interIter = memResMap_.find(tag);
      83           15 :     if (interIter != memResMap_.end()) {
      84           16 :         memResMap_.erase(tag);
      85              :     }
      86           16 :     lock.unlock();
      87              : 
      88           15 :     return HCCL_SUCCESS;
      89           15 : }
      90              : 
      91            0 : void WorkSpaceMem::DestroyMemResource()
      92              : {
      93            0 :     HCCL_INFO("[Destroy][MemResource]Destroy workspace all mem");
      94            0 :     std::unique_lock<std::mutex> lock(memResMutex_);
      95            0 :     memResMap_.clear();
      96            0 :     lock.unlock();
      97            0 : }
      98              : 
      99           95 : bool WorkSpaceMem::IsExist(const std::string& tag)
     100              : {
     101           95 :     std::unique_lock<std::mutex> lock(memResMutex_);
     102           95 :     auto interIter = memResMap_.find(tag);
     103          190 :     return interIter != memResMap_.end();
     104           95 : }
     105              : } // namespace hccl