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