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              : #ifndef HCCLV2_TOKENINFO_MANAGER_H
      12              : #define HCCLV2_TOKENINFO_MANAGER_H
      13              : 
      14              : #include <mutex>
      15              : #include <vector>
      16              : #include <unordered_map>
      17              : #include <hccl/hccl_types.h>
      18              : #include "buffer_key.h"
      19              : #include "orion_adapter_hccp.h"
      20              : 
      21              : namespace Hccl {
      22              : 
      23              : using TokenInfo = std::pair<TokenIdHandle, uint32_t>;
      24              : using BufKeyVecIndex = u32;
      25              : constexpr BufKeyVecIndex INVALID_BUF_KEY_VEC_INDEX = UINT32_MAX;
      26              : 
      27              : class TokenRefMap
      28              : {
      29              : public:
      30              :     using Iterator = std::unordered_map<BufKeyVecIndex, TokenInfo>::iterator;
      31              : 
      32           30 :     Iterator begin()
      33              :     {
      34           30 :         return data_.begin();
      35              :     }
      36              : 
      37           34 :     Iterator end()
      38              :     {
      39           34 :         return data_.end();
      40              :     }
      41              : 
      42          141 :     u32 insert(BufKeyVecIndex key, const TokenInfo &value)
      43              :     {
      44          141 :         if (has(key)) {
      45            0 :             ref_[key]++;
      46              :         } else {
      47          141 :             data_.insert(std::make_pair(key, value));
      48          141 :             ref_[key] = 1;
      49              :         }
      50          141 :         return count(key);
      51              :     }
      52              : 
      53           47 :     u32 eraseAndGet(BufKeyVecIndex key, TokenInfo &erasedValue)
      54              :     {
      55           47 :         u32 refCount = count(key);
      56           47 :         if (refCount > 1) {
      57            0 :             ref_[key]--;
      58           47 :         } else if (refCount == 1) {
      59           47 :             erasedValue = std::move(data_[key]);
      60           47 :             data_.erase(key);
      61           47 :             ref_.erase(key);
      62              :         }
      63           47 :         return count(key);
      64              :     }
      65              : 
      66           30 :     void clear()
      67              :     {
      68           30 :         data_.clear();
      69           30 :         ref_.clear();
      70           30 :     }
      71              : 
      72          705 :     bool has(BufKeyVecIndex key)
      73              :     {
      74          705 :         return data_.find(key) != data_.end();
      75              :     }
      76              : 
      77          376 :     u32 count(BufKeyVecIndex key)
      78              :     {
      79          376 :         return has(key) ? ref_[key] : 0;
      80              :     }
      81              : 
      82          141 :     TokenInfo &operator[](BufKeyVecIndex key)
      83              :     {
      84          141 :         return data_[key];
      85              :     }
      86              : 
      87              : private:
      88              :     std::unordered_map<BufKeyVecIndex, TokenInfo> data_;
      89              :     std::unordered_map<BufKeyVecIndex, u32> ref_;
      90              : };
      91              : 
      92              : class TokenInfoManager
      93              : {
      94              : public:
      95           43 :     TokenInfoManager(u32 devId, RdmaHandle rdmahandle) : devId_(devId), rdmahandle_(rdmahandle)
      96              :     {
      97           43 :     }
      98              :     
      99              :     TokenInfo GetTokenInfo(const BufferKey<uintptr_t, u64> &bufKey);
     100              :     void PutTokenInfo(const BufferKey<uintptr_t, u64> &bufKey, TokenIdHandle tokenIdHandle);
     101              : 
     102              :     void Destroy();
     103              : 
     104              : private:
     105              :     u32        devId_;
     106              :     RdmaHandle rdmahandle_;
     107              :     std::mutex tokenInfoMgrMutex_;
     108              : 
     109              :     TokenRefMap tokenRefMap_;
     110              :     std::unordered_map<u32, vector<vector<BufferKey<uintptr_t, u64>>>> bufferKeysMap_; // <devId, BufKeyVecIndex, vector<BufferKey>>
     111              :     std::unordered_map<TokenIdHandle, BufKeyVecIndex> tokenIdToIndex_;
     112              : 
     113              :     BufKeyVecIndex GetBufferVecIndex(const BufferKey<uintptr_t, u64> &inputBufKey);
     114              : };
     115              : 
     116              : bool HasIntersect(const vector<BufferKey<uintptr_t, u64>> &bufKeys, const BufferKey<uintptr_t, u64> &inputBufKey);
     117              : 
     118              : } // namespace Hccl
     119              : 
     120              : #endif // HCCLV2_TOKENINFO_MANAGER_H