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 CCU_CHANNELCTX_POOLS_H
      12              : #define CCU_CHANNELCTX_POOLS_H
      13              : 
      14              : #include <vector>
      15              : #include <unordered_map>
      16              : #include <mutex>
      17              : 
      18              : #include "ccu_jetty_.h"
      19              : #include "hash_utils.h"
      20              : #include "ip_address.h"
      21              : #include "virtual_topo.h"
      22              : 
      23              : namespace hcomm {
      24              : 
      25              : // 管理着有限的硬件资源：ChannelCtx与jetty
      26              : class CcuChannelCtxPool final {
      27              : public:
      28              :     explicit CcuChannelCtxPool(int32_t devLogicId);
      29              :     ~CcuChannelCtxPool();
      30              : 
      31              :     HcclResult PrepareCreate(const std::vector<Hccl::LinkData>& links, uint32_t sqSize = 0);
      32              :     using CcuChannelCtx = std::pair<CcuChannelInfo, std::vector<CcuJetty*>>;
      33              :     HcclResult GetChannelCtx(const Hccl::LinkData& link, CcuChannelCtx& channelCtx) const;
      34              :     HcclResult GetCcuChannelCtxById(const std::pair<uint8_t, uint32_t>& key, CcuChannelCtx& ctx);
      35              :     // Channel 销毁时归还 channel ctx / jetty ctx / wqeBB 资源：
      36              :     // - V1：所属 batch 立即变空，逐 channel 归还设备层空闲池并销毁 batch；
      37              :     // - V2：槽位压回 batch 的可复用列表供后续创建复用，整组无活跃 channel 时
      38              :     //       再逐 channel 归还设备层空闲池（与设备层 useCnt 组粒度语义一致）。
      39              :     HcclResult ReleaseChannel(const Hccl::LinkData& link);
      40              : 
      41              : private:
      42              :     struct ResIdHash {
      43          234 :         std::size_t operator()(const std::pair<uint8_t, uint32_t>& p) const
      44              :         {
      45          234 :             return Hccl::HashCombine({p.first, p.second});
      46              :         }
      47              :     };
      48              : 
      49              :     using CcuJettyPtr = CcuJetty*;
      50              :     using BatchKey = Hccl::IpAddress; // srcIpAddress;
      51              :     using ResIdkey = std::pair<uint8_t, uint32_t>;
      52              :     using ChannelIdKey = ResIdkey;
      53              :     using JettyIdKey = ResIdkey;
      54              : 
      55              :     // 平台层每次调用CcuAllocChannels可能提供多个ccu channel，且不同srcIp的jetty不能复用
      56              :     // 故以srcIp为粒度，多次调用接口，每次接口结果定义为一个批次资源
      57              :     struct ResourceBatch { // 记录该批次申请到的所有channel资源信息
      58              :         BatchKey key;
      59              :         std::vector<ChannelIdKey> channelIdKeys;
      60              :         std::vector<ChannelIdKey> availableChannelIdKeys;
      61              :         std::unordered_map<JettyIdKey, std::unique_ptr<CcuJetty>, ResIdHash> jettys;
      62              : 
      63           15 :         ResourceBatch(const BatchKey& batchKey) : key(batchKey) {};
      64              :         HcclResult Init(const std::vector<CcuChannelInfo>& channelInfos);
      65              :     };
      66              : 
      67              : private:
      68              :     HcclResult GetAvailableBatch(const BatchKey& batchKey, ResourceBatch*& batchPtr, uint32_t sqSize);
      69              :     bool FindAvailableBatch(const BatchKey& batchKey, ResourceBatch*& batchPtr) const;
      70              :     HcclResult CreateAndSaveNewBatch(
      71              :         const BatchKey& batchKey, const std::vector<CcuChannelInfo> channelInfos, ResourceBatch*& batchPtr);
      72              :     HcclResult ReleaseConfirmedChannelRes();
      73              :     ResourceBatch* FindBatchByChannelId(const ChannelIdKey& key) const;
      74              :     HcclResult ReleaseBatchIfIdle(ResourceBatch* batch);
      75              :     // 从 batchMap_ 移除并销毁该 batch（锁内调用）：设备层 CcuReleaseChannel 与
      76              :     // ~CcuJetty（RaCtxQpDestroy）在 pool 锁内执行，牺牲并发流畅，换取"释放先于
      77              :     // 同一 pool 的后续申请"，资源紧俏场景下保证释放的资源可被立即复用。
      78              :     void RemoveBatch(ResourceBatch* batch);
      79              : 
      80              : private:
      81              :     int32_t devLogicId_{0};
      82              :     bool isReleased_{true};
      83              : 
      84              :     // 保护以下所有 map/batch，PrepareCreate/GetChannelCtx/ReleaseChannel 并发安全
      85              :     mutable std::mutex mtx_;
      86              :     // 各资源申请记录，当前按SrcIpAddr粒度申请和管理
      87              :     std::unordered_map<BatchKey, std::vector<std::unique_ptr<ResourceBatch>>> batchMap_;
      88              :     // 各link已分配的channel资源Id信息
      89              :     std::unordered_map<Hccl::LinkData, ChannelIdKey> allocatedChannelIdMap_;
      90              :     // 全部已申请的channel资源信息，资源申请成功后将要记录到该map中
      91              :     std::unordered_map<ChannelIdKey, CcuChannelCtx, ResIdHash> channelJettyInfoMap_;
      92              :     // 以die粒度记录已分配channel资源, index: dieId
      93              :     std::unordered_map<uint8_t, uint32_t> usedChannelCntMap_;
      94              :     // 记录channel与对端rank的映射关系, index: (die, channelId)
      95              :     std::unordered_map<ChannelIdKey, Hccl::RankId, ResIdHash> channelRemoteRankIdMap_;
      96              :     // channel -> 所属 batch 反向索引，ReleaseChannel 定位 batch 用
      97              :     std::unordered_map<ChannelIdKey, ResourceBatch*, ResIdHash> channelToBatch_{};
      98              : };
      99              : 
     100              : } // namespace hcomm
     101              : 
     102              : #endif // CCU_CHANNELCTX_POOLS_H