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          521 : WorkSpaceMem::WorkSpaceMem() = default;
      17              : 
      18          520 : 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("[Set][MemResource]set mem resource success, tag[%s] ptr[%llu] max size[%llu]",
      35              :             tag.c_str(), std::hash<void *>{}(ptr), maxSize);
      36          250 :     }
      37          250 :     return HCCL_SUCCESS;
      38              : }
      39              : 
      40           24 : void* WorkSpaceMem::AllocMem(const std::string &tag, u64 size)
      41              : {
      42              :     /* 此处可能会与并发，加锁 */
      43           24 :     std::unique_lock<std::mutex> lock(memResMutex_);
      44              : 
      45           23 :     auto interIter = memResMap_.find(tag);
      46           24 :     if (interIter == memResMap_.end()) {
      47           12 :         HCCL_ERROR("[Alloc][Mem] Can't find tag, tag[%s], size[%llu]", tag.c_str(), size);
      48           12 :         return nullptr;
      49              :     }
      50           11 :     if (interIter->second.curPtr == nullptr) {
      51            0 :         HCCL_ERROR("[Alloc][Mem] Current resource(curPtr) is null, tag[%s], size[%llu]", tag.c_str(), size);
      52            0 :         return nullptr;
      53              :     }
      54              : 
      55              :     // 判断size 是否超出最大值
      56           11 :     if (interIter->second.maxSize < (interIter->second.totalSize + size)) {
      57            0 :         HCCL_ERROR("[Alloc][Mem] Current resource size is not enough, tag[%s], size[%llu], "\
      58              :             "maxSize[%llu], totalSize[%llu]", tag.c_str(), size,
      59              :             interIter->second.maxSize, interIter->second.totalSize);
      60            0 :         return nullptr;
      61              :     }
      62              : 
      63           11 :     void* tempPtr = interIter->second.curPtr;
      64           10 :     interIter->second.curPtr = reinterpret_cast<void *>(static_cast<char *>(interIter->second.curPtr) + size);
      65           10 :     interIter->second.totalSize += size;
      66              : 
      67           10 :     HCCL_INFO("[Alloc][Mem]Alloc mem success, tag[%s] size[%llu], totalSize[%llu], maxSize[%llu]",
      68              :         tag.c_str(), size, interIter->second.totalSize, interIter->second.maxSize);
      69           11 :     lock.unlock();
      70           12 :     return tempPtr;
      71           24 : }
      72              : 
      73           16 : HcclResult WorkSpaceMem::DestroyMemResource(const std::string &tag)
      74              : {
      75           16 :     HCCL_INFO("[Destroy][MemResource]Destroy workspace mem, tag[%s]", tag.c_str());
      76              : 
      77           16 :     std::unique_lock<std::mutex> lock(memResMutex_);
      78              : 
      79           14 :     auto interIter = memResMap_.find(tag);
      80           14 :     if (interIter != memResMap_.end()) {
      81           15 :         memResMap_.erase(tag);
      82              :     }
      83           13 :     lock.unlock();
      84              : 
      85           11 :     return HCCL_SUCCESS;
      86           14 : }
      87              : 
      88            0 : void WorkSpaceMem::DestroyMemResource()
      89              : {
      90            0 :     HCCL_INFO("[Destroy][MemResource]Destroy workspace all mem");
      91            0 :     std::unique_lock<std::mutex> lock(memResMutex_);
      92            0 :     memResMap_.clear();
      93            0 :     lock.unlock();
      94            0 : }
      95              : 
      96           95 : bool WorkSpaceMem::IsExist(const std::string &tag)
      97              : {
      98           95 :     std::unique_lock<std::mutex> lock(memResMutex_);
      99           95 :     auto interIter = memResMap_.find(tag);
     100          190 :     return interIter != memResMap_.end();
     101           95 : }
     102              : }  // namespace hccl