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