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 <unordered_set>
16 : #include <memory>
17 : #include <mutex>
18 : #include "hccl/hccl_res.h"
19 : #include "hccl_independent_common.h"
20 : #include "aicpu_ts_thread.h"
21 : #include "cpu_ts_thread.h"
22 : #include "log.h"
23 : #include "manager_common.h"
24 :
25 : namespace hccl {
26 :
27 : class ThreadMgr {
28 : public:
29 : ThreadMgr(
30 : uint32_t threadNum, uint32_t notifyNumPerThread, std::string commId, aclrtBinHandle binHandle,
31 : const ManagerCallbacks& callbacks);
32 : ~ThreadMgr();
33 : HcclResult HcclThreadAcquire(
34 : CommEngine engine, uint32_t threadNum, ThreadType type, const ThreadConfig* config, ThreadHandle* threads,
35 : std::vector<uint32_t>& threadId);
36 : HcclResult HcclThreadAcquireV2(
37 : CommEngine engine, uint32_t threadNum, ThreadType type, const ThreadConfig* config, ThreadHandle* threads,
38 : std::vector<uint32_t>& threadId);
39 : HcclResult
40 : HcclThreadAcquireWithStream(CommEngine engine, rtStream_t stream, uint32_t notifyNum, ThreadHandle* thread);
41 : HcclResult HcclGetNotifyNumInThread(ThreadHandle thread, uint32_t* notifyNum);
42 : HcclResult HcclThreadExportToCommEngine(
43 : uint32_t threadNum, const ThreadHandle* threads, CommEngine dstCommEngine, ThreadHandle* exportedThreads);
44 : HcclResult HcclThreadResGetInfo(ThreadHandle thread, ThreadResType resType, uint32_t infoLen, void** info);
45 : HcclResult
46 : HcclDedicatedThreadAcquire(HcclDedicatedThreadType useType, uint32_t notifyNumPerThread, ThreadHandle* thread);
47 : HcclResult RegisterOrderLaunchThread(ThreadHandle thread);
48 21 : u32 GetThreadNum() const { return threadNum_; }
49 : u32 GetNotifyNumPerThread() const { return notifyNumPerThread_; }
50 :
51 : private:
52 : uint64_t GetMaxNotifyTotal();
53 : HcclResult CheckNotifyNum(CommEngine engine, uint32_t threadNum, uint32_t notifyNumPerThread);
54 : HcclResult CheckThreadNum(CommEngine engine, uint32_t threadNum, uint32_t notifyNumPerThread);
55 : HcclResult SupplementNotify(CommEngine engine, std::vector<std::shared_ptr<Thread>>& needSupplementThread);
56 : HcclResult SupplementNotify(CommEngine engine, uint32_t threadNum, ThreadType type, const ThreadConfig* config);
57 : HcclResult SupplementThread(
58 : CommEngine engine, std::vector<std::shared_ptr<Thread>>& newThreads,
59 : std::unique_ptr<ThreadHandle[]>& hostHandle);
60 : HcclResult
61 : SupplementThread(CommEngine engine, uint32_t supplementThreadNum, ThreadType type, const ThreadConfig* config);
62 : HcclResult
63 : ThreadExportToCommEngineCpu(uint32_t threadNum, const ThreadHandle* threads, ThreadHandle* exportedThreads);
64 : HcclResult ThreadExportToCommEngineAicpu(
65 : uint32_t threadNum, const ThreadHandle* threads, CommEngine dstCommEngine, ThreadHandle* exportedThreads);
66 : HcclResult GetExportedThread(
67 : const ThreadHandle threadHandle, CommEngine commEngine, Thread*& exportedThread,
68 : std::shared_ptr<Thread>& threadOut);
69 : HcclResult ExportHostThreadsToAicpu(
70 : std::vector<std::shared_ptr<Thread>>& hostThreads, const std::vector<u32>& index, const ThreadHandle* threads,
71 : CommEngine dstCommEngine, ThreadHandle* exportedThreads);
72 : HcclResult ExportOrderLaunchThreadsToAicpu(
73 : std::vector<std::shared_ptr<Thread>>& orderLaunchHostThreads, const std::vector<u32>& orderLaunchIndex,
74 : const ThreadHandle* threads, CommEngine dstCommEngine, ThreadHandle* exportedThreads);
75 : HcclResult CreateAndInitThreads(
76 : CommEngine engine, StreamType streamType, NotifyLoadType notifyLoadType, uint32_t threadNum,
77 : const ThreadConfig* config, std::vector<std::shared_ptr<Thread>>& newThreads);
78 : HcclResult AssignThreadHandles(
79 : CommEngine engine, std::vector<std::shared_ptr<Thread>>& newThreads, ThreadHandle* threads,
80 : std::unique_ptr<ThreadHandle[]>& hostHandle);
81 : HcclResult StoreThreadsAndBuildHandleMap(
82 : CommEngine engine, std::vector<std::shared_ptr<Thread>>& newThreads,
83 : std::unique_ptr<ThreadHandle[]>& hostHandle);
84 : HcclResult
85 : HcclUnfoldThreadAcquire(HcclDedicatedThreadType useType, uint32_t notifyNumPerThread, ThreadHandle* thread);
86 :
87 : u32 threadNum_ = 0;
88 : u32 notifyNumPerThread_ = 0;
89 : std::string commId_;
90 : aclrtBinHandle binHandle_;
91 :
92 : u64 usedNotifyNum_ = 0;
93 : std::mutex threadMutex_;
94 : std::vector<std::shared_ptr<Thread>> threads_;
95 :
96 : std::mutex mainThreadMutex_;
97 : std::map<rtStream_t, std::shared_ptr<Thread>> mainThread_;
98 :
99 : std::mutex engineToThreadMutex_;
100 : std::map<std::pair<CommEngine, ThreadType>, std::vector<std::shared_ptr<Thread>>> engineToThreadsMap_;
101 :
102 : std::mutex threadMapMutex_;
103 : std::unordered_map<ThreadHandle, ThreadHandle>
104 : threadHandleOthersToCpu_; // 其他引擎上的ThreadHandle与CPU_TS上的ThreadHandle的映射
105 : std::unordered_map<ThreadHandle, ThreadHandle> hostToDeviceThreadHandle_;
106 : ManagerCallbacks callbacks_;
107 :
108 : std::mutex threadhandleToThreadMutex_;
109 : std::unordered_map<ThreadHandle, std::shared_ptr<Thread>> threadMap_;
110 :
111 : std::mutex dedicatedThreadMutex_;
112 : std::unordered_map<HcclDedicatedThreadType, ThreadHandle> dedicatedThreadMap_;
113 :
114 : std::unordered_set<ThreadHandle>
115 : orderLaunchThreads_; // 保序流（非所有权，资源由 OrderLaunchThreadMgr 管理，同 comm 操作串行无需加锁）
116 : };
117 : } // namespace hccl
118 : #endif
|
