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 COLL_COMM_AICPU_MGR_H
      12              : #define COLL_COMM_AICPU_MGR_H
      13              : 
      14              : #include "coll_comm_aicpu.h"
      15              : #include <shared_mutex>
      16              : #include <unordered_map>
      17              : #include <vector>
      18              : #include <string>
      19              : #include <utility>
      20              : 
      21              : class CollCommAicpuMgr {
      22              : public:
      23              :     static CollCommAicpuMgr& GetInstance();
      24              : 
      25              :     // 通信域初始化 (原 AicpuIndOpCommInit)
      26              :     HcclResult InitComm(CommAicpuParam* commAicpuParam);
      27              : 
      28              :     // 通信域注册表操作
      29              :     CollCommAicpu* AcquireCommForUse(const std::string& group); // 获取并标记使用中 (原 AicpuGetCommMgrbyGroup)
      30              :     HcclResult
      31              :     AcquireAndCreateComm(const std::string& group, CollCommAicpu** outComm); // 创建或获取通信域（不标记使用中）
      32              :     void ReleaseComm(const std::string& group);               // 释放使用标记 (原 AicpuReleaseCommMgrbyGroup)
      33              :     CollCommAicpu* GetCurrentComm(const std::string& group);  // 获取当前线程通信域 (原 AicpuGetComm)
      34              :     CollCommAicpu* FindCommByGroup(const std::string& group); // 从 map 按 group 查找
      35              :     HcclResult DestroyComm(const std::string& group);
      36              :     HcclResult GetAllComms(std::vector<std::pair<std::string, CollCommAicpu*>>& aicpuCommInfo);
      37              :     std::shared_mutex& GetMutex();
      38            0 :     CollCommAicpu* GetCurrentComm() { return currentComm_; }
      39              : 
      40              :     // 全局环境初始化 (原 CollCommAicpu::InitIndopEnv / InitBackGroundThread)
      41              :     void InitIndopEnv(CommAicpuParam* commAicpuParam);
      42              :     void InitBackGroundThread(u32 devId);
      43              : 
      44              : private:
      45            5 :     CollCommAicpuMgr() = default;
      46            5 :     ~CollCommAicpuMgr() = default;
      47              :     CollCommAicpuMgr(const CollCommAicpuMgr&) = delete;
      48              :     CollCommAicpuMgr& operator=(const CollCommAicpuMgr&) = delete;
      49              : 
      50              :     struct CommEntry {
      51              :         std::unique_ptr<CollCommAicpu> comm;
      52              :         bool isUsed{false};
      53              :     };
      54              : 
      55              :     std::shared_mutex commMapMutex_;
      56              :     std::unordered_map<std::string, CommEntry> commMap_;
      57              :     static thread_local CollCommAicpu* currentComm_;
      58              : };
      59              : 
      60              : #endif // COLL_COMM_AICPU_MGR_H