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 "tokenInfo_manager.h"
      12              : 
      13              : namespace Hccl {
      14              : 
      15          141 : TokenInfo TokenInfoManager::GetTokenInfo(const BufferKey<uintptr_t, u64> &bufKey)
      16              : {
      17          141 :     std::lock_guard<std::mutex> lock(tokenInfoMgrMutex_);
      18              : 
      19          141 :     BufKeyVecIndex index = GetBufferVecIndex(bufKey);
      20          141 :     if (!tokenRefMap_.has(index)) {
      21          141 :         TokenInfo tokenInfo = RaUbAllocTokenIdHandle(rdmahandle_);
      22          141 :         tokenRefMap_.insert(index, tokenInfo);
      23          141 :         tokenIdToIndex_[tokenInfo.first] = index;
      24              :     } else {
      25            0 :         tokenRefMap_.insert(index, tokenRefMap_[index]);
      26              :     }
      27          423 :     HCCL_INFO("[TokenInfoManager::%s] rdmahandle[%p] index[%u] refCount[%u]",
      28              :               __func__, rdmahandle_, index, tokenRefMap_.count(index));
      29          282 :     return tokenRefMap_[index];
      30          141 : }
      31              : 
      32          111 : void TokenInfoManager::PutTokenInfo(const BufferKey<uintptr_t, u64> &bufKey, TokenIdHandle tokenIdHandle)
      33              : {
      34          111 :     std::lock_guard<std::mutex> lock(tokenInfoMgrMutex_);
      35              : 
      36          111 :     auto mapIt = tokenIdToIndex_.find(tokenIdHandle);
      37          111 :     if (mapIt == tokenIdToIndex_.end()) {
      38          192 :         HCCL_WARNING("[TokenInfoManager::%s] tokenIdHandle[0x%llx] not found, skip", __func__, tokenIdHandle);
      39           64 :         return;
      40              :     }
      41           47 :     BufKeyVecIndex index = mapIt->second;
      42              : 
      43           47 :     if (!tokenRefMap_.has(index)) {
      44            0 :         HCCL_WARNING("[TokenInfoManager::%s] index[%u] not in tokenRefMap, skip", __func__, index);
      45            0 :         return;
      46              :     }
      47              : 
      48           47 :     auto &bufferKeys = bufferKeysMap_[devId_];
      49           47 :     if (index < bufferKeys.size()) {
      50           47 :         auto it = std::find_if(bufferKeys[index].begin(), bufferKeys[index].end(),
      51           47 :             [bufKey](const auto &curBufKey) { return bufKey == curBufKey; });
      52           47 :         if (it != bufferKeys[index].end()) {
      53           47 :             bufferKeys[index].erase(it);
      54              :         }
      55              :     }
      56              : 
      57           47 :     TokenInfo erasedInfo{};
      58           47 :     u32 remain = tokenRefMap_.eraseAndGet(index, erasedInfo);
      59          141 :     HCCL_DEBUG("[TokenInfoManager::%s] rdmahandle[%p] index[%u] remain[%u]",
      60              :               __func__, rdmahandle_, index, remain);
      61              : 
      62           47 :     if (remain == 0) {
      63           47 :         DECTOR_TRY_CATCH("token id handle free", RaUbFreeTokenIdHandle(rdmahandle_, erasedInfo.first));
      64           47 :         tokenIdToIndex_.erase(erasedInfo.first);
      65           47 :         if (index < bufferKeys.size()) {
      66           47 :             bufferKeys[index].clear();
      67              :         }
      68              :     }
      69          111 : }
      70              : 
      71         5091 : bool HasIntersect(const vector<BufferKey<uintptr_t, u64>> &bufKeys, const BufferKey<uintptr_t, u64> &inputBufKey)
      72              : {
      73              :     // 遍历bufKeys, 若bufKeys中存在和inputBufKey相交的bufKey则返回true, 否则fasle
      74         5091 :     auto it = std::find_if(bufKeys.begin(), bufKeys.end(),
      75         5050 :                            [inputBufKey](const auto &curBufKeyInVec) {
      76         5050 :                                 return inputBufKey.IsIntersect(curBufKeyInVec);
      77              :                            });
      78         5091 :     return it != bufKeys.end();
      79              : }
      80              : 
      81          141 : BufKeyVecIndex TokenInfoManager::GetBufferVecIndex(const BufferKey<uintptr_t, u64> &inputBufKey)
      82              : {
      83          141 :     auto &bufferKeys = bufferKeysMap_[devId_];
      84          141 :     u32 size = bufferKeys.size();
      85          141 :     auto it = std::find_if(bufferKeys.begin(), bufferKeys.end(),
      86         5091 :                            [inputBufKey](const auto &unOverlapBufVec) {
      87         5091 :                                 return !HasIntersect(unOverlapBufVec, inputBufKey);
      88              :                             });
      89              :     // 若不存在一组bufferKey与inputBufKey不相交, 需要申请新的token, 返回新的索引
      90          141 :     u32 idx = std::distance(bufferKeys.begin(), it);
      91          141 :     if (idx == size) {
      92          300 :         bufferKeys.push_back(vector<BufferKey<uintptr_t, u64>>{inputBufKey});
      93              :     } else {
      94              :         // 若存在一组bufferKey与inputBufKey不相交则返回该组索引, 然后将inputBufKey插入
      95           41 :         it->push_back(inputBufKey);
      96              :     }
      97              : 
      98          423 :     HCCL_INFO("[TokenInfoManager::%s] idx[%u] size[%u]", __func__, idx, size);
      99          141 :     return idx;
     100              : }
     101              : 
     102           30 : void TokenInfoManager::Destroy()
     103              : {
     104           86 :     HCCL_INFO("[TokenInfoManager::%s] rdmahandle[%p]", __func__, rdmahandle_);
     105           34 :     for (auto it = tokenRefMap_.begin(); it != tokenRefMap_.end(); ++it) {
     106            4 :         DECTOR_TRY_CATCH("token id handle destroy",
     107              :                          RaUbFreeTokenIdHandle(rdmahandle_, it->second.first));
     108              :     }
     109           30 :     tokenRefMap_.clear();
     110           30 :     tokenIdToIndex_.clear();
     111           30 : }
     112              : 
     113              : } // namespace Hccl