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 : #include "hccl_team_mgr.h"
12 :
13 : #include <algorithm>
14 : #include <cstdint>
15 : #include <memory>
16 : #include <mutex>
17 : #include <utility>
18 :
19 : #include "adapter_rts_common.h"
20 : #include "log.h"
21 :
22 : namespace hccl {
23 :
24 682 : HcclTeamMgr& HcclTeamMgr::GetInstance()
25 : {
26 : static std::once_flag instanceFlag;
27 : static HcclTeamMgr* instance = nullptr;
28 682 : std::call_once(instanceFlag, [&] {
29 6 : instance = new HcclTeamMgr();
30 6 : });
31 682 : return *instance;
32 : }
33 :
34 30 : HcclResult HcclTeamMgr::RegisterWorldTeam(
35 : HcommTeamHandle worldTeam, CollComm* collComm, void* syncMemPtr, uint64_t syncMemSize, const uint32_t* rankIds,
36 : uint32_t rankNum)
37 : {
38 30 : if (worldTeam == nullptr || collComm == nullptr || rankIds == nullptr) {
39 0 : HCCL_ERROR(
40 : "[HcclTeamMgr][%s] invalid param, worldTeam[%p] collComm[%p] rankIds[%p]", __func__, worldTeam, collComm,
41 : rankIds);
42 0 : return HCCL_E_PTR;
43 : }
44 30 : std::unique_lock<std::shared_mutex> lock(mutex_);
45 30 : auto worldIt = teamMap_.find(worldTeam);
46 30 : if (worldIt != teamMap_.end()) {
47 0 : HCCL_ERROR("[HcclTeamMgr][%s] worldTeam[%p] already registered", __func__, worldTeam);
48 0 : return HCCL_E_PARA;
49 : }
50 30 : TeamEntry& entry = teamMap_[worldTeam];
51 30 : entry.collComm = collComm;
52 30 : entry.worldTeam = nullptr;
53 30 : entry.syncMemPtr = syncMemPtr;
54 30 : entry.syncMemSize = syncMemSize;
55 30 : entry.syncMemHandle = nullptr;
56 30 : entry.syncMemTag.clear();
57 30 : entry.windows.clear();
58 30 : entry.rankIds.assign(rankIds, rankIds + rankNum);
59 30 : return HCCL_SUCCESS;
60 30 : }
61 :
62 3 : HcclResult HcclTeamMgr::RegisterSubTeam(
63 : HcommTeamHandle worldTeam, HcommTeamHandle subTeam, void* syncMemPtr, uint64_t syncMemSize, const uint32_t* rankIds,
64 : uint32_t rankNum)
65 : {
66 3 : if (worldTeam == nullptr || subTeam == nullptr || rankIds == nullptr) {
67 0 : HCCL_ERROR(
68 : "[HcclTeamMgr][%s] invalid param, worldTeam[%p] subTeam[%p] rankIds[%p]", __func__, worldTeam, subTeam,
69 : rankIds);
70 0 : return HCCL_E_PTR;
71 : }
72 3 : std::unique_lock<std::shared_mutex> lock(mutex_);
73 3 : auto worldIt = teamMap_.find(worldTeam);
74 3 : if (worldIt == teamMap_.end()) {
75 0 : HCCL_ERROR("[HcclTeamMgr][%s] worldTeam[%p] not registered", __func__, worldTeam);
76 0 : return HCCL_E_PARA;
77 : }
78 3 : if (teamMap_.find(subTeam) != teamMap_.end()) {
79 0 : HCCL_ERROR("[HcclTeamMgr][%s] subTeam[%p] already registered", __func__, subTeam);
80 0 : return HCCL_E_PARA;
81 : }
82 3 : CollComm* collComm = worldIt->second.collComm;
83 3 : TeamEntry& entry = teamMap_[subTeam];
84 3 : entry.collComm = collComm;
85 3 : entry.worldTeam = worldTeam;
86 3 : entry.syncMemPtr = syncMemPtr;
87 3 : entry.syncMemSize = syncMemSize;
88 3 : entry.syncMemHandle = nullptr;
89 3 : entry.syncMemTag.clear();
90 3 : entry.windows.clear();
91 3 : entry.rankIds.assign(rankIds, rankIds + rankNum);
92 3 : return HCCL_SUCCESS;
93 3 : }
94 :
95 129 : void HcclTeamMgr::UnregisterTeam(HcommTeamHandle team)
96 : {
97 129 : if (team == nullptr) {
98 0 : return;
99 : }
100 129 : void* syncMemPtr = nullptr;
101 : {
102 129 : std::unique_lock<std::shared_mutex> lock(mutex_);
103 129 : auto it = teamMap_.find(team);
104 129 : if (it == teamMap_.end()) {
105 97 : return;
106 : }
107 32 : syncMemPtr = it->second.syncMemPtr;
108 32 : teamMap_.erase(it);
109 129 : }
110 32 : if (syncMemPtr != nullptr) {
111 32 : (void)hrtFree(syncMemPtr);
112 : }
113 : }
114 :
115 33 : CollComm* HcclTeamMgr::FindCollComm(HcommTeamHandle team)
116 : {
117 33 : CHK_PRT_RET(team == nullptr, HCCL_ERROR("[HcclTeamMgr][%s] team[%p] is invalid", __func__, team), nullptr);
118 33 : std::shared_lock<std::shared_mutex> lock(mutex_);
119 33 : auto it = teamMap_.find(team);
120 33 : if (it == teamMap_.end()) {
121 3 : HCCL_ERROR("[HcclTeamMgr][%s] team[%p] not registered", __func__, team);
122 3 : return nullptr;
123 : }
124 30 : return it->second.collComm;
125 33 : }
126 :
127 76 : HcommTeamHandle HcclTeamMgr::FindWorldTeam(HcommTeamHandle team)
128 : {
129 76 : CHK_PRT_RET(team == nullptr, HCCL_ERROR("[HcclTeamMgr][%s] team[%p] is invalid", __func__, team), nullptr);
130 76 : std::shared_lock<std::shared_mutex> lock(mutex_);
131 76 : auto it = teamMap_.find(team);
132 76 : if (it == teamMap_.end()) {
133 5 : HCCL_ERROR("[HcclTeamMgr][%s] team[%p] not registered", __func__, team);
134 5 : return nullptr;
135 : }
136 : // world team 自身 worldTeam 字段为 nullptr，返回自身
137 71 : return (it->second.worldTeam != nullptr) ? it->second.worldTeam : team;
138 76 : }
139 :
140 28 : std::vector<uint32_t> HcclTeamMgr::GetRankIds(HcommTeamHandle team)
141 : {
142 28 : std::vector<uint32_t> result;
143 28 : if (team == nullptr) {
144 0 : return result;
145 : }
146 28 : std::shared_lock<std::shared_mutex> lock(mutex_);
147 28 : auto it = teamMap_.find(team);
148 28 : if (it == teamMap_.end()) {
149 0 : return result;
150 : }
151 28 : result = it->second.rankIds;
152 28 : return result;
153 28 : }
154 :
155 18 : void* HcclTeamMgr::GetSyncMemPtr(HcommTeamHandle team)
156 : {
157 18 : CHK_PRT_RET(team == nullptr, HCCL_ERROR("[HcclTeamMgr][%s] team[%p] is invalid", __func__, team), nullptr);
158 18 : std::shared_lock<std::shared_mutex> lock(mutex_);
159 18 : auto it = teamMap_.find(team);
160 18 : if (it == teamMap_.end()) {
161 0 : HCCL_ERROR("[HcclTeamMgr][%s] team[%p] not registered", __func__, team);
162 0 : return nullptr;
163 : }
164 18 : return it->second.syncMemPtr;
165 18 : }
166 :
167 18 : uint64_t HcclTeamMgr::GetSyncMemSize(HcommTeamHandle team)
168 : {
169 18 : CHK_PRT_RET(team == nullptr, HCCL_ERROR("[HcclTeamMgr][%s] team[%p] is invalid", __func__, team), 0);
170 18 : std::shared_lock<std::shared_mutex> lock(mutex_);
171 18 : auto it = teamMap_.find(team);
172 18 : if (it == teamMap_.end()) {
173 0 : HCCL_ERROR("[HcclTeamMgr][%s] team[%p] not registered", __func__, team);
174 0 : return 0;
175 : }
176 18 : return it->second.syncMemSize;
177 18 : }
178 :
179 11 : void HcclTeamMgr::SetTeamSyncMemHandle(HcommTeamHandle team, HcclMemHandle handle, const std::string& tag)
180 : {
181 11 : if (team == nullptr) {
182 0 : HCCL_ERROR("[HcclTeamMgr][%s] team[%p] is invalid", __func__, team);
183 0 : return;
184 : }
185 11 : std::unique_lock<std::shared_mutex> lock(mutex_);
186 11 : auto it = teamMap_.find(team);
187 11 : if (it == teamMap_.end()) {
188 0 : HCCL_ERROR("[HcclTeamMgr][%s] team[%p] not registered", __func__, team);
189 0 : return;
190 : }
191 11 : it->second.syncMemHandle = handle;
192 11 : it->second.syncMemTag = tag;
193 11 : }
194 :
195 11 : HcclMemHandle HcclTeamMgr::GetTeamSyncMemHandle(HcommTeamHandle team)
196 : {
197 11 : CHK_PRT_RET(team == nullptr, HCCL_ERROR("[HcclTeamMgr][%s] team[%p] is invalid", __func__, team), nullptr);
198 11 : std::shared_lock<std::shared_mutex> lock(mutex_);
199 11 : auto it = teamMap_.find(team);
200 11 : if (it == teamMap_.end()) {
201 0 : HCCL_ERROR("[HcclTeamMgr][%s] team[%p] not registered", __func__, team);
202 0 : return nullptr;
203 : }
204 11 : return it->second.syncMemHandle;
205 11 : }
206 :
207 11 : std::string HcclTeamMgr::GetTeamSyncMemTag(HcommTeamHandle team)
208 : {
209 11 : CHK_PRT_RET(team == nullptr, HCCL_ERROR("[HcclTeamMgr][%s] team[%p] is invalid", __func__, team), std::string());
210 11 : std::shared_lock<std::shared_mutex> lock(mutex_);
211 11 : auto it = teamMap_.find(team);
212 11 : if (it == teamMap_.end()) {
213 0 : HCCL_ERROR("[HcclTeamMgr][%s] team[%p] not registered", __func__, team);
214 0 : return std::string();
215 : }
216 11 : return it->second.syncMemTag;
217 11 : }
218 :
219 21 : bool HcclTeamMgr::FindReusableWindow(HcommTeamHandle worldTeam, const CommMem& localMem, HcommWindowHandle& window)
220 : {
221 21 : if (worldTeam == nullptr) {
222 0 : return false;
223 : }
224 21 : std::unique_lock<std::shared_mutex> lock(mutex_);
225 21 : auto it = teamMap_.find(worldTeam);
226 21 : if (it == teamMap_.end()) {
227 0 : return false;
228 : }
229 : // windows 按 registeredLocalMem.addr 升序
230 : // upper_bound 找第一个 start > requestStart 的，反向遍历找包含请求区间的 window（请求是已注册 window 的子集）。
231 21 : const uintptr_t requestStart = reinterpret_cast<uintptr_t>(localMem.addr);
232 21 : const uintptr_t requestEnd = requestStart + localMem.size;
233 21 : auto& windows = it->second.windows;
234 : auto upper
235 21 : = std::upper_bound(windows.begin(), windows.end(), requestStart, [](uintptr_t addr, const WindowInfo& win) {
236 12 : return addr < reinterpret_cast<uintptr_t>(win.registeredLocalMem.addr);
237 : });
238 30 : for (auto wIt = upper; wIt != windows.begin();) {
239 14 : --wIt;
240 14 : if (wIt->handle == nullptr) {
241 0 : continue;
242 : }
243 14 : const uintptr_t winStart = reinterpret_cast<uintptr_t>(wIt->registeredLocalMem.addr);
244 14 : const uintptr_t winEnd = winStart + wIt->registeredLocalMem.size;
245 : // 请求区间是已注册 window 的子集（含精确匹配）→ 复用
246 14 : if (requestStart >= winStart && requestEnd <= winEnd) {
247 5 : window = wIt->handle;
248 : // 精确匹配（addr+size 完全相同）→ 已有记录，无需新增
249 5 : if (requestStart == winStart && requestEnd == winEnd) {
250 2 : return true;
251 : }
252 : // 子集包含但非精确匹配 → 复用 parent handle，为请求的 localMem 新增一条 WindowInfo，
253 : // 使后续 BindWindowsAndSyncMem 的 self 槽填入请求的实际地址（而非 parent 的地址范围）。
254 3 : WindowInfo info;
255 3 : info.handle = window;
256 3 : info.registeredLocalMem = localMem;
257 3 : auto insertIt = std::upper_bound(
258 6 : windows.begin(), windows.end(), requestStart, [](uintptr_t addr, const WindowInfo& win) {
259 6 : return addr < reinterpret_cast<uintptr_t>(win.registeredLocalMem.addr);
260 : });
261 3 : windows.insert(insertIt, std::move(info));
262 3 : return true;
263 3 : }
264 : }
265 16 : return false;
266 21 : }
267 :
268 19 : void HcclTeamMgr::AddWorldTeamWindow(
269 : HcommTeamHandle worldTeam, HcommWindowHandle window, const CommMem& localMem, HcclMemHandle localMemHandle,
270 : const std::string& localMemTag)
271 : {
272 19 : if (worldTeam == nullptr || window == nullptr) {
273 0 : HCCL_ERROR("[HcclTeamMgr][%s] worldTeam[%p] or window[%p] is invalid", __func__, worldTeam, window);
274 0 : return;
275 : }
276 19 : std::unique_lock<std::shared_mutex> lock(mutex_);
277 19 : auto it = teamMap_.find(worldTeam);
278 19 : if (it == teamMap_.end()) {
279 0 : HCCL_ERROR("[HcclTeamMgr][%s] worldTeam[%p] not registered", __func__, worldTeam);
280 0 : return;
281 : }
282 19 : WindowInfo info;
283 19 : info.handle = window;
284 19 : info.registeredLocalMem = localMem;
285 19 : info.localMemHandle = localMemHandle;
286 19 : info.localMemTag = localMemTag;
287 : // 按 registeredLocalMem.addr 升序插入，维护 windows 有序性，供 FindReusableWindow 二分查找
288 19 : const uintptr_t newStart = reinterpret_cast<uintptr_t>(localMem.addr);
289 38 : auto insertIt = std::upper_bound(
290 38 : it->second.windows.begin(), it->second.windows.end(), newStart, [](uintptr_t addr, const WindowInfo& win) {
291 4 : return addr < reinterpret_cast<uintptr_t>(win.registeredLocalMem.addr);
292 : });
293 19 : it->second.windows.insert(insertIt, std::move(info));
294 19 : }
295 :
296 43 : std::vector<WindowInfo> HcclTeamMgr::GetWorldTeamWindows(HcommTeamHandle worldTeam)
297 : {
298 43 : std::vector<WindowInfo> result;
299 43 : if (worldTeam == nullptr) {
300 0 : return result;
301 : }
302 43 : std::shared_lock<std::shared_mutex> lock(mutex_);
303 43 : auto it = teamMap_.find(worldTeam);
304 43 : if (it == teamMap_.end()) {
305 0 : return result;
306 : }
307 43 : result = it->second.windows;
308 43 : return result;
309 43 : }
310 :
311 11 : std::vector<HcclMemHandle> HcclTeamMgr::CollectPendingMemHandles(HcommTeamHandle worldTeam, HcommTeamHandle team)
312 : {
313 11 : std::vector<HcclMemHandle> result;
314 11 : if (worldTeam == nullptr || team == nullptr) {
315 0 : return result;
316 : }
317 11 : std::unique_lock<std::shared_mutex> lock(mutex_);
318 11 : auto itTeam = teamMap_.find(team);
319 11 : if (itTeam == teamMap_.end()) {
320 0 : return result;
321 : }
322 11 : auto itWorldTeam = teamMap_.find(worldTeam);
323 11 : if (itWorldTeam == teamMap_.end()) {
324 0 : return result;
325 : }
326 : // syncMemHandle：首次注册且未交换时收集
327 11 : if (!itTeam->second.syncMemExchanged && itTeam->second.syncMemHandle != nullptr) {
328 11 : result.push_back(itTeam->second.syncMemHandle);
329 11 : itTeam->second.syncMemExchanged = true;
330 : }
331 : // window localMemHandle：未交换的收集并标记
332 19 : for (auto& win : itWorldTeam->second.windows) {
333 8 : if (!win.exchanged && win.localMemHandle != nullptr) {
334 8 : result.push_back(win.localMemHandle);
335 8 : win.exchanged = true;
336 : }
337 : }
338 11 : return result;
339 11 : }
340 :
341 3 : void HcclTeamMgr::RemoveWorldTeamWindow(HcommTeamHandle worldTeam, HcommWindowHandle window)
342 : {
343 3 : if (worldTeam == nullptr || window == nullptr) {
344 2 : return;
345 : }
346 3 : std::unique_lock<std::shared_mutex> lock(mutex_);
347 3 : auto it = teamMap_.find(worldTeam);
348 3 : if (it == teamMap_.end()) {
349 0 : return;
350 : }
351 3 : auto& wins = it->second.windows;
352 3 : for (auto wIt = wins.begin(); wIt != wins.end(); ++wIt) {
353 2 : if (wIt->handle == window) {
354 2 : wins.erase(wIt);
355 2 : return;
356 : }
357 : }
358 3 : }
359 :
360 29 : std::vector<HcommTeamHandle> HcclTeamMgr::GetSubTeams(HcommTeamHandle worldTeam)
361 : {
362 29 : std::vector<HcommTeamHandle> result;
363 29 : if (worldTeam == nullptr) {
364 0 : return result;
365 : }
366 29 : std::shared_lock<std::shared_mutex> lock(mutex_);
367 : // subTeam 的 worldTeam 字段指向其父 worldTeam；worldTeam 自身该字段为 nullptr。
368 59 : for (const auto& pair : teamMap_) {
369 30 : if (pair.second.worldTeam == worldTeam) {
370 1 : result.push_back(pair.first);
371 : }
372 : }
373 29 : return result;
374 29 : }
375 :
376 188 : void HcclTeamMgr::ClearByCollComm(CollComm* collComm)
377 : {
378 188 : if (collComm == nullptr) {
379 0 : return;
380 : }
381 : // 锁内收集并 erase，锁外销毁，避免持 L2 锁调 L3 接口（锁嵌套/死锁风险）
382 188 : std::vector<TeamCleanupInfo> cleanupInfos = CollectTeamCleanupInfo(collComm);
383 188 : ExecuteTeamCleanup(cleanupInfos);
384 188 : }
385 :
386 188 : std::vector<HcclTeamMgr::TeamCleanupInfo> HcclTeamMgr::CollectTeamCleanupInfo(CollComm* collComm)
387 : {
388 188 : std::vector<TeamCleanupInfo> cleanupInfos;
389 188 : std::unique_lock<std::shared_mutex> lock(mutex_);
390 189 : for (auto it = teamMap_.begin(); it != teamMap_.end();) {
391 1 : if (it->second.collComm != collComm) {
392 0 : ++it;
393 0 : continue;
394 : }
395 1 : TeamCleanupInfo info;
396 1 : info.handle = it->first;
397 1 : info.syncMemPtr = it->second.syncMemPtr;
398 : // window 归 worldTeam 所有（entry.worldTeam==nullptr 表示自身是 worldTeam）
399 1 : if (it->second.worldTeam == nullptr) {
400 2 : for (const auto& win : it->second.windows) {
401 1 : if (win.handle != nullptr) {
402 1 : info.windows.push_back(win.handle);
403 : }
404 : }
405 : }
406 1 : cleanupInfos.push_back(std::move(info));
407 1 : it = teamMap_.erase(it);
408 1 : }
409 188 : return cleanupInfos;
410 188 : }
411 :
412 188 : void HcclTeamMgr::ExecuteTeamCleanup(const std::vector<TeamCleanupInfo>& cleanupInfos)
413 : {
414 : // 锁外依次销毁：window（L3 devWindow/devMems）→ team（L3 device 资源）→ syncMem（L2 本地内存）
415 189 : for (const auto& info : cleanupInfos) {
416 2 : for (HcommWindowHandle win : info.windows) {
417 1 : (void)HcommTeamWindowDeregister(info.handle, win);
418 : }
419 1 : (void)HcommTeamDestroy(info.handle);
420 1 : if (info.syncMemPtr != nullptr) {
421 1 : (void)hrtFree(info.syncMemPtr);
422 : }
423 : }
424 188 : }
425 :
426 : } // namespace hccl
|
