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          232 : TokenInfo TokenInfoManager::GetTokenInfo(const BufferKey<uintptr_t, u64> &bufKey)
      16              : {
      17          232 :     std::lock_guard<std::mutex> lock(tokenInfoMgrMutex_);
      18              : 
      19          232 :     BufKeyVecIndex index = GetBufferVecIndex(bufKey);
      20          232 :     auto iter = tokenInfoMap_.find(index);
      21          232 :     if (iter == tokenInfoMap_.end()) {
      22           56 :         TokenInfo tokenInfo = RaUbAllocTokenIdHandle(rdmahandle_);
      23           56 :         tokenInfoMap_.emplace(index, tokenInfo);
      24              :     }
      25          696 :     HCCL_INFO("[TokenInfoManager::%s] rdmahandle[%p] index[%u]", __func__, rdmahandle_, index);
      26          464 :     return tokenInfoMap_[index];
      27          232 : }
      28              : 
      29          998 : bool HasIntersect(const vector<BufferKey<uintptr_t, u64>> &bufKeys, const BufferKey<uintptr_t, u64> &inputBufKey)
      30              : {
      31              :     // 遍历bufKeys, 若bufKeys中存在和inputBufKey相交的bufKey则返回true, 否则fasle
      32          998 :     auto it = std::find_if(bufKeys.begin(), bufKeys.end(), 
      33         9824 :                            [inputBufKey](const auto &curBufKeyInVec) { 
      34         9824 :                                 return inputBufKey.IsIntersect(curBufKeyInVec);
      35              :                            });
      36          998 :     return it != bufKeys.end();
      37              : }
      38              : 
      39          232 : BufKeyVecIndex TokenInfoManager::GetBufferVecIndex(const BufferKey<uintptr_t, u64> &inputBufKey)
      40              : {
      41          232 :     auto &bufferKeys = bufferKeysMap_[devId_];
      42          232 :     u32 size = bufferKeys.size();
      43          232 :     auto it = std::find_if(bufferKeys.begin(), bufferKeys.end(), 
      44          998 :                            [inputBufKey](const auto &unOverlapBufVec) { 
      45          998 :                                 return !HasIntersect(unOverlapBufVec, inputBufKey); 
      46              :                             });
      47              : 
      48              :     // 若不存在一组bufferKey与inputBufKey不相交, 需要申请新的token, 返回新的索引
      49          232 :     u32 idx = std::distance(bufferKeys.begin(), it);
      50          232 :     if (idx == size) {
      51          168 :         bufferKeys.push_back(vector<BufferKey<uintptr_t, u64>>{inputBufKey});
      52              :     } else {
      53              :         // 若存在一组bufferKey与inputBufKey不相交则返回该组索引, 然后将inputBufKey插入
      54          176 :         it->push_back(inputBufKey);
      55              :     }
      56              : 
      57          696 :     HCCL_INFO("[TokenInfoManager::%s] idx[%u] size[%u]", __func__, idx, size);
      58          232 :     return idx;
      59              : }
      60              : 
      61           30 : void TokenInfoManager::Destroy()
      62              : {
      63           86 :     HCCL_INFO("[TokenInfoManager::%s] rdmahandle[%p]", __func__, rdmahandle_);
      64           39 :     for (auto &tokenInfo : tokenInfoMap_) {
      65            9 :         DECTOR_TRY_CATCH("token id handle destroy",
      66              :                          RaUbFreeTokenIdHandle(rdmahandle_, tokenInfo.second.first));
      67              :     }
      68           30 : }
      69              : 
      70              : } // namespace Hccl