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 "comm_mem_manager.h"
      12              : namespace hccl {
      13              : 
      14          234 : void CommMemMgr::CommSetHcclBufferManager(CCLBufferManager &bufferManager)
      15              : {
      16          234 :     bufferManager_ = &bufferManager;
      17          234 : }
      18              : 
      19            7 : HcclResult CommMemMgr::GetHcclBuffer(CommBuffer *buffer)
      20              : {
      21            7 :     CHK_PTR_NULL(buffer);
      22            7 :     CHK_PTR_NULL(bufferManager_);
      23            7 :     std::lock_guard<std::mutex> lock(bufferMutex_);
      24            7 :     void* temp = nullptr;
      25            7 :     uint64_t tempSize = 0;
      26            7 :     HcclResult ret = bufferManager_->GetIndependentOpCCLbuffer(temp, tempSize);
      27            7 :     CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[GetHcclBuffer] GetHcclBuffer failed"), ret);
      28            7 :     buffer->addr = temp;
      29            7 :     buffer->size = tempSize;
      30            7 :     return HCCL_SUCCESS;
      31            7 : }
      32              : 
      33              : // 绑定：opTag -> 句柄（幂等）
      34            0 : HcclResult CommMemMgr::CommRegMem(const std::string& memTag, const HcclMem& mem, HcclRegMemAttr attr,
      35              :     void **memHandle)
      36              : {
      37            0 :     CHK_PRT_RET(memHandle == nullptr, HCCL_ERROR("[CommRegMem] memHandle is null. tag[%s]", memTag.c_str()), HCCL_E_PARA);
      38            0 :     CHK_PRT_RET(mem.addr == nullptr || mem.size == 0, HCCL_ERROR("[CommRegMem] invalid mem. addr[%p] size[%llu]",
      39              :         mem.addr, mem.size), HCCL_E_PARA);
      40              : 
      41              :     // 组装句柄（仅域内管理，无进程级注册）
      42            0 :     Handle h;
      43            0 :     EXCEPTION_CATCH(h = std::make_shared<HcclMemoryHandle>(), return HCCL_E_PTR);
      44            0 :     h->addr    = mem.addr;
      45            0 :     h->size    = static_cast<uint64_t>(mem.size);
      46            0 :     h->memType = static_cast<HcclMemType>(mem.type);
      47            0 :     h->attr    = attr;
      48              : 
      49            0 :     const auto key = MakeKey(mem.addr, static_cast<size_t>(mem.size));
      50              : 
      51            0 :     std::lock_guard<std::mutex> addLock(memMutex_);
      52            0 :     auto& reg = tagRegs_[memTag];
      53              : 
      54              :     // 同tag内做区间冲突/幂等复用
      55            0 :     auto res = reg.table.Add(key, h);
      56            0 :     if (!res.second) {
      57              :         // 只能用 Find 的返回值来判定：
      58              :         // - 等于(全集命中)：Find(key).first == true（允许，Add 内已 ref）
      59              :         // - 子集/超集/交集：Find(key).first 可能为 true(子) 或 false(交/超/空)，但都属于冲突！
      60            0 :         auto f = reg.table.Find(key);
      61            0 :         if (!f.first || !(f.second && f.second->addr == mem.addr && f.second->size == mem.size)) {
      62            0 :             HCCL_ERROR("[CommRegMem] overlap in tag[%s], key=%s", memTag.c_str(), key.ToString().c_str());
      63            0 :             return HCCL_E_PARA;
      64              :         }
      65              :         // HcclRegMemAttr不同时更新
      66            0 :         if (f.second->attr.value != attr.value) {
      67            0 :             HCCL_WARNING("[CommRegMem] inconsistent attr for same mem. tag[%s]", memTag.c_str());
      68            0 :             f.second->attr.value = attr.value;
      69              :         }
      70              :         // 复用已有句柄：直接用 Find 返回的 buffer，避免解引用 res.first（可能是 end()）
      71            0 :         h = f.second;
      72            0 :     }
      73              : 
      74              :     // 幂等加入绑定列表（同memHandle不重复）
      75            0 :     auto& vec = opBindings_[memTag];
      76            0 :     bool exists = std::any_of(vec.begin(), vec.end(),
      77            0 :         [&h](const Handle& x){ return x && (x.get() == h.get()); });
      78            0 :     if (!exists) vec.emplace_back(h);
      79              : 
      80            0 :     *memHandle = h.get();
      81            0 :     HCCL_INFO("[CommRegMem] ok. tag[%s] memHandle[%p] size[%llu]", memTag.c_str(), *memHandle, h->size);
      82            0 :     return HCCL_SUCCESS;
      83            0 : }
      84              : 
      85              : // 解绑：在该通信域实例内，移除“指定算子(memTag)”下的该句柄
      86            0 : HcclResult CommMemMgr::CommUnregMem(const std::string& memTag, const void* memHandle)
      87              : {
      88            0 :     CHK_PRT_RET(memHandle == nullptr, HCCL_ERROR("[CommUnregMem] memHandle is null"), HCCL_E_PARA);
      89            0 :     CHK_PRT_RET(memTag.empty(), HCCL_ERROR("[CommUnregMem] memTag is null or empty"), HCCL_E_PARA);
      90              : 
      91            0 :     std::lock_guard<std::mutex> addLock(memMutex_);
      92              : 
      93            0 :     auto itTag = opBindings_.find(memTag);
      94            0 :     CHK_PRT_RET(itTag == opBindings_.end(),
      95              :         HCCL_WARNING("[CommUnregMem] tag[%s] not found in bindings", memTag.c_str()), HCCL_E_NOT_FOUND);
      96              : 
      97            0 :     auto &vec = itTag->second;                // vector<Handle> under this tag
      98            0 :     auto &reg = tagRegs_[itTag->first];       // TagRegistry for this tag
      99            0 :     size_t unboundCount = 0;  // 本次解绑命中的句柄个数（即便 Del 未真正擦除也计数）
     100            0 :     size_t erasedCount  = 0;  // RmaBufferMgr::Del 返回 true 的次数（ref 归零而“擦除”）
     101              : 
     102            0 :     vec.erase(std::remove_if(vec.begin(), vec.end(),
     103            0 :         [&](const Handle &h) {
     104            0 :             if (!h || h.get() != memHandle) return false;
     105            0 :             const auto key = MakeKey(h->addr, static_cast<size_t>(h->size));
     106              :             try {
     107            0 :                 if (reg.table.Del(key)) {
     108            0 :                     ++erasedCount;            // 该 key 的引用归零并从表中移除
     109              :                 }
     110            0 :             } catch (const std::out_of_range &) {
     111            0 :                 HCCL_ERROR("[CommUnregMem] tag[%s] key not found on Del (maybe already removed)", itTag->first.c_str());
     112            0 :             }
     113              : 
     114            0 :             ++unboundCount;                   // 从绑定列表移除，无论 Del 是否真正擦除
     115            0 :             return true;                      // erase-remove：删除该 handle
     116              :         }),
     117            0 :         vec.end());
     118              : 
     119              :     // 若该 tag 已无绑定，可按需清理映射条目（以及空表）
     120            0 :     if (vec.empty()) {
     121            0 :         opBindings_.erase(itTag);
     122            0 :         if (reg.table.size() == 0) {
     123            0 :             tagRegs_.erase(std::string(memTag));
     124              :         }
     125              :     }
     126              : 
     127            0 :     CHK_PRT_RET(unboundCount == 0,
     128              :         HCCL_WARNING("[CommUnregMem] tag[%s] memHandle[%p] not found", memTag.c_str(), memHandle), HCCL_E_NOT_FOUND);
     129              : 
     130            0 :     HCCL_INFO("[CommUnregMem] tag[%s] memHandle[%p] unbound=%zu, erased=%zu",
     131              :               memTag.c_str(), memHandle, unboundCount, erasedCount);
     132            0 :     return HCCL_SUCCESS;
     133            0 : }
     134              : 
     135            0 : HcclResult CommMemMgr::CommGetLocalRegMemByTag(const std::string &tag,
     136              :                                                std::vector<HcclMem> &memVec)
     137              : {
     138            0 :     std::lock_guard<std::mutex> lock(memMutex_);
     139            0 :     auto it = opBindings_.find(tag);
     140            0 :     if (it == opBindings_.end()) {
     141            0 :         HCCL_INFO("[CommMemMgr] tag[%s] key not found", tag.c_str());
     142            0 :         return HCCL_SUCCESS;
     143              :     }
     144              : 
     145            0 :     const auto &vec = it->second;
     146            0 :     memVec.reserve(vec.size());
     147            0 :     for (const auto &handle : vec) {
     148              :         HcclMem mem;
     149            0 :         mem.addr = handle->addr;
     150            0 :         mem.size = handle->size;
     151            0 :         mem.type = handle->memType;
     152            0 :         memVec.push_back(mem);
     153              :     }
     154            0 :     return HCCL_SUCCESS;
     155            0 : }
     156              : }