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 THREAD_MANAGER_H
      12              : #define THREAD_MANAGER_H
      13              : #include <string>
      14              : #include <unordered_map>
      15              : #include <memory>
      16              : #include <mutex>
      17              : #include "hccl/hccl_res.h"
      18              : #include "hccl_independent_common.h"
      19              : #include "aicpu_ts_thread.h"
      20              : #include "cpu_ts_thread.h"
      21              : #include "log.h"
      22              : #include "manager_common.h"
      23              : 
      24              : namespace hccl {
      25              : 
      26              : class ThreadMgr {
      27              : public:
      28              :     ThreadMgr(uint32_t threadNum, uint32_t notifyNumPerThread, std::string commId, aclrtBinHandle binHandle, const ManagerCallbacks& callbacks);
      29              :     ~ThreadMgr();
      30              :     HcclResult HcclThreadAcquire(CommEngine engine, uint32_t threadNum, ThreadType type,
      31              :         const ThreadConfig *config, ThreadHandle *threads, std::vector<uint32_t> &threadId);
      32              :     HcclResult HcclThreadAcquireV2(CommEngine engine, uint32_t threadNum, ThreadType type,
      33              :         const ThreadConfig *config, ThreadHandle *threads, std::vector<uint32_t> &threadId);
      34              :     HcclResult HcclThreadAcquireWithStream(CommEngine engine,
      35              :         rtStream_t stream, uint32_t notifyNum, ThreadHandle *thread);
      36              :     HcclResult HcclGetNotifyNumInThread(ThreadHandle thread, uint32_t *notifyNum);
      37              :     HcclResult HcclThreadExportToCommEngine(uint32_t threadNum, const ThreadHandle *threads, CommEngine dstCommEngine, ThreadHandle *exportedThreads);
      38              :     HcclResult HcclThreadResGetInfo(ThreadHandle thread, ThreadResType resType, uint32_t infoLen, void **info);
      39              :     HcclResult HcclDedicatedThreadAcquire( HcclDedicatedThreadType useType, uint32_t notifyNumPerThread, ThreadHandle *thread);
      40           21 :     u32 GetThreadNum() const { return threadNum_; }
      41              :     u32 GetNotifyNumPerThread() const { return notifyNumPerThread_; }
      42              : 
      43              : private:
      44              :     uint64_t GetMaxNotifyTotal();
      45              :     HcclResult CheckNotifyNum(CommEngine engine, uint32_t threadNum, uint32_t notifyNumPerThread);
      46              :     HcclResult CheckThreadNum(CommEngine engine, uint32_t threadNum, uint32_t notifyNumPerThread);
      47              :     HcclResult SupplementNotify(CommEngine engine, std::vector<std::shared_ptr<Thread>> &needSupplementThread);
      48              :     HcclResult SupplementNotify(CommEngine engine, uint32_t threadNum, ThreadType type, const ThreadConfig *config);
      49              :     HcclResult SupplementThread(CommEngine engine, std::vector<std::shared_ptr<Thread>> &newThreads, std::unique_ptr<ThreadHandle[]> &hostHandle);
      50              :     HcclResult SupplementThread(CommEngine engine, uint32_t supplementThreadNum, ThreadType type, const ThreadConfig *config);
      51              :     HcclResult ThreadExportToCommEngineCpu(uint32_t threadNum, const ThreadHandle *threads, ThreadHandle *exportedThreads);
      52              :     HcclResult ThreadExportToCommEngineAicpu(uint32_t threadNum, const ThreadHandle *threads, CommEngine dstCommEngine, ThreadHandle *exportedThreads);
      53              :     HcclResult GetExportedThread(const ThreadHandle threadHandle, CommEngine commEngine, Thread *&exportedThread, std::shared_ptr<Thread> &threadOut);
      54              :     HcclResult CreateAndInitThreads(CommEngine engine, StreamType streamType,
      55              :         NotifyLoadType notifyLoadType, uint32_t threadNum, const ThreadConfig *config,
      56              :         std::vector<std::shared_ptr<Thread>> &newThreads);
      57              :     HcclResult AssignThreadHandles(CommEngine engine,
      58              :         std::vector<std::shared_ptr<Thread>> &newThreads, ThreadHandle *threads,
      59              :         std::unique_ptr<ThreadHandle[]> &hostHandle);
      60              :     HcclResult StoreThreadsAndBuildHandleMap(CommEngine engine,
      61              :         std::vector<std::shared_ptr<Thread>> &newThreads,
      62              :         std::unique_ptr<ThreadHandle[]> &hostHandle);
      63              :     HcclResult HcclUnfoldThreadAcquire(HcclDedicatedThreadType useType, uint32_t notifyNumPerThread, ThreadHandle *thread);
      64              : 
      65              :     u32 threadNum_ = 0;
      66              :     u32 notifyNumPerThread_ = 0;
      67              :     std::string commId_;
      68              :     aclrtBinHandle binHandle_;
      69              : 
      70              :     u64 usedNotifyNum_ = 0;
      71              :     std::mutex threadMutex_;
      72              :     std::vector<std::shared_ptr<Thread>> threads_;
      73              : 
      74              :     std::mutex mainThreadMutex_;
      75              :     std::map<rtStream_t, std::shared_ptr<Thread>> mainThread_;
      76              : 
      77              :     std::mutex engineToThreadMutex_;
      78              :     std::map<std::pair<CommEngine, ThreadType>, std::vector<std::shared_ptr<Thread>>> engineToThreadsMap_;
      79              : 
      80              :     std::mutex threadMapMutex_;
      81              :     std::unordered_map<ThreadHandle, ThreadHandle> threadHandleOthersToCpu_; // 其他引擎上的ThreadHandle与CPU_TS上的ThreadHandle的映射
      82              :     std::unordered_map<ThreadHandle, ThreadHandle> hostToDeviceThreadHandle_;
      83              :     ManagerCallbacks callbacks_;
      84              : 
      85              :     std::mutex threadhandleToThreadMutex_;
      86              :     std::unordered_map<ThreadHandle, std::shared_ptr<Thread>> threadMap_;
      87              : 
      88              :     std::mutex dedicatedThreadMutex_; 
      89              :     std::unordered_map<HcclDedicatedThreadType, ThreadHandle> dedicatedThreadMap_;
      90              : };
      91              : }
      92              : #endif