Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2026 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 HCOMM_TEAM_MGR_H
      12              : #define HCOMM_TEAM_MGR_H
      13              : 
      14              : #include <cstdint>
      15              : #include <memory>
      16              : #include <mutex>
      17              : #include <shared_mutex>
      18              : #include <unordered_map>
      19              : #include <vector>
      20              : 
      21              : #include "hcomm_res.h"
      22              : #include "hcomm_res_defs.h"
      23              : #include "hcomm_team.h"
      24              : #include "hcomm_team_defs.h"
      25              : #include "hcomm_team_entity_defs.h"
      26              : 
      27              : namespace hcomm {
      28              : // HcommTeam / HcommWindow 结构体的 ABI 头部常量
      29              : constexpr uint32_t HCOMM_TEAM_MAGIC_WORD = 0x0f0f0f20U;
      30              : constexpr uint32_t HCOMM_TEAM_VERSION = 1U;
      31              : constexpr uint32_t HCOMM_WINDOW_MAGIC_WORD = 0x0f0f0f21U;
      32              : constexpr uint32_t HCOMM_WINDOW_VERSION = 1U;
      33              : 
      34              : struct WindowEntry {
      35              :     HcommWindow hostWindow{};
      36              :     HcommWindowHandle devWindow{nullptr};
      37              :     void* devMems{nullptr};
      38              :     CommMem* hostMems{nullptr};
      39              :     HcommTeamHandle teamHandle{nullptr};
      40              : };
      41              : 
      42              : struct TeamEntry {
      43              :     HcommTeam hostTeam{};
      44              :     HcommTeamHandle devTeam{nullptr};
      45              :     void* devWorldTeamIds{nullptr};
      46              :     void* devChannelNums{nullptr};
      47              :     void* devChannels{nullptr};
      48              :     void* devRemoteMems{nullptr};
      49              : 
      50              :     uint32_t* hostWorldTeamIds{nullptr};
      51              :     std::vector<uint32_t> hostChannelNums;
      52              :     std::vector<std::vector<uint64_t>> channelsList;
      53              : 
      54              :     CommMem* hostRemoteMems{nullptr};
      55              : 
      56              :     HcommTeamHandle worldTeamHandle{nullptr};
      57              :     bool isSubTeam{false};
      58              : 
      59              :     HcommTeamSyncMemRequirement syncMemReq{};
      60              :     uint64_t syncMemSize{0};
      61              : };
      62              : 
      63              : class HcommTeamMgr {
      64              : public:
      65              :     static HcommTeamMgr& GetInstance();
      66              : 
      67              :     HcommResult TeamCreate(
      68              :         HcommTeamHandle worldTeam, const HcommTeamCreateDesc* desc, HcommTeamHandle* team, uint64_t* outSyncMemSize);
      69              :     HcommResult TeamDestroy(HcommTeamHandle team);
      70              :     HcommResult WindowRegister(HcommTeamHandle team, HcommWindowHandle* handle);
      71              :     HcommResult BindWindow(HcommTeamHandle team, HcommWindowHandle handle, const HcommTeamWindowDesc* desc);
      72              :     HcommResult WindowDeregister(HcommTeamHandle team, HcommWindowHandle handle);
      73              :     HcommResult BindChannels(HcommTeamHandle team, const HcommTeamBindChannelsDesc* desc);
      74              :     HcommResult BindSyncMem(HcommTeamHandle team, const HcommTeamBindSyncMemDesc* desc);
      75              :     HcommResult GetNetLayer(HcommTeamHandle team, uint32_t* netLayer);
      76              : 
      77              : private:
      78            1 :     HcommTeamMgr() = default;
      79              :     ~HcommTeamMgr();
      80              :     HcommTeamMgr(const HcommTeamMgr&) = delete;
      81              :     HcommTeamMgr& operator=(const HcommTeamMgr&) = delete;
      82              : 
      83              :     TeamEntry* FindTeamByHandleLocked(HcommTeamHandle handle);
      84              :     WindowEntry* FindWindowByHandleLocked(HcommWindowHandle handle);
      85              :     HcommResult SyncTeamToDevice(TeamEntry* entry);
      86              :     HcommResult SyncWindowToDevice(WindowEntry* entry);
      87              :     HcommResult AllocAndCopyWorldTeamIds(TeamEntry* entry, const uint32_t* src, uint32_t memberNum);
      88              :     HcommResult AllocAndCopyChannels(TeamEntry* entry);
      89              :     HcommResult AllocChannelEntities(TeamEntry* entry);
      90              :     HcommResult AllocChannelNumsArray(TeamEntry* entry);
      91              :     HcommResult AllocAndCopyRemoteMems(TeamEntry* entry, const CommMem* src, uint32_t memberNum);
      92              :     // 首次分配 hostRemoteMems（calloc+memcpy）与 devRemoteMems（hrtMalloc+hrtMemSyncCopy），并写入 syncMem 字段
      93              :     HcommResult AllocRemoteMems(TeamEntry* entry, const CommMem* src, uint32_t memberNum);
      94              :     // 重复 bind：校验维度一致后更新 hostRemoteMems（memcpy）并重新 sync 到 devRemoteMems
      95              :     HcommResult UpdateRemoteMems(TeamEntry* entry, const CommMem* src, uint32_t memberNum);
      96              :     HcommResult ValidateSubTeam(TeamEntry* worldEntry, const HcommTeamCreateDesc* desc);
      97              :     void InitTeamEntry(TeamEntry* entry, const HcommTeamCreateDesc* desc, HcommTeamHandle worldTeam);
      98              :     HcommResult AllocAndSyncTeam(TeamEntry* entry, const HcommTeamCreateDesc* desc);
      99              :     HcommResult AllocAndCopyWindowMems(WindowEntry* winEntry, uint64_t memberNum, const CommMem* src);
     100              :     HcommResult MergeWindowMems(WindowEntry* winEntry, uint64_t memberNum, const CommMem* src);
     101              :     void MergeChannelLists(
     102              :         TeamEntry* entry, const HcommTeamBindChannelsDesc* desc, std::vector<std::vector<uint64_t>>& newChannels);
     103              :     void FreeTeamResources(TeamEntry* entry);
     104              :     void FreeWindowResources(WindowEntry* entry);
     105              :     void FreeDeviceChannels(TeamEntry* entry);
     106              : 
     107              :     std::shared_mutex teamsRwMutex_; // 保护 teams_
     108              :     std::unordered_map<HcommTeamHandle, std::unique_ptr<TeamEntry>> teams_;
     109              :     std::shared_mutex windowsRwMutex_; // 保护 windows_
     110              :     std::unordered_map<HcommWindowHandle, std::unique_ptr<WindowEntry>> windows_;
     111              :     std::shared_mutex windowToTeamRwMutex_; // 保护 windowToTeamMap_
     112              :     std::unordered_map<HcommWindowHandle, HcommTeamHandle> windowToTeamMap_;
     113              : };
     114              : 
     115              : } // namespace hcomm
     116              : 
     117              : #endif // HCOMM_TEAM_MGR_H