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 : #include "mirror_task_manager.h"
11 :
12 : namespace Hccl {
13 :
14 389 : MirrorTaskManager::MirrorTaskManager(u32 devId, GlobalMirrorTasks *globalMirrorTasks, bool devUsed)
15 389 : : devId_(devId), globalMirrorTasks_(globalMirrorTasks), devUsed_(devUsed)
16 : {
17 389 : }
18 :
19 262 : void MirrorTaskManager::RegFullyCallBack(std::function<void()> callBack)
20 : {
21 262 : fullyCallBack_ = callBack;
22 262 : return;
23 : }
24 :
25 29 : QueueType MirrorTaskManager::GetQueueType() const
26 : {
27 29 : if (currDfxOpInfo_ == nullptr) {
28 15 : HCCL_WARNING("[MirrorTaskManager][%s]currDfxOpInfo_ is nullptr, return default Circular_Queue!", __func__);
29 5 : return QueueType::Circular_Queue;
30 : }
31 24 : QueueType queueType = QueueType::Vector_Queue;
32 :
33 24 : if (devUsed_ || isStaticGraphMode_ || (opMode_ == OpMode::OPBASE)) {
34 24 : queueType = QueueType::Circular_Queue;
35 : }
36 24 : return queueType;
37 : }
38 :
39 10107 : void MirrorTaskManager::AddTaskInfo(std::unique_ptr<TaskInfo> &&taskInfo)
40 : {
41 10107 : if (UNLIKELY(taskInfo == nullptr)) {
42 1 : THROW<InternalException>(
43 3 : StringFormat("MirrorTaskManager::AddTaskInfo taskInfo is nullptr"));
44 : }
45 10106 : bool needCallback = false;
46 10106 : std::unique_lock<std::mutex> lock(profMutex);
47 10106 : if (taskInfo->dfxOpInfo_ == nullptr) {
48 54 : taskInfo->dfxOpInfo_ = currDfxOpInfo_;
49 : }
50 :
51 10106 : auto emplaceResult = streamQueues_.emplace(taskInfo->streamId_, MirrorStreamQueueEntry{nullptr, QueueType::Vector_Queue, 0});
52 10106 : MirrorStreamQueueEntry *entryPtr = &emplaceResult.first->second;
53 10106 : if (emplaceResult.second) {
54 28 : entryPtr->queueType = GetQueueType();
55 28 : entryPtr->queue = &(globalMirrorTasks_->CreateQueue(devId_, taskInfo->streamId_, entryPtr->queueType));
56 : }
57 10106 : if (UNLIKELY(entryPtr->taskNum == entryPtr->queue->Capacity())) {
58 4 : needCallback = true;
59 4 : entryPtr->taskNum = 0;
60 : }
61 :
62 10106 : if (needCallback && fullyCallBack_ != nullptr) {
63 4 : lock.unlock();
64 4 : fullyCallBack_();
65 4 : lock.lock();
66 4 : auto queueIt = streamQueues_.find(taskInfo->streamId_);
67 4 : if (queueIt == streamQueues_.end()) {
68 0 : THROW<InternalException>(
69 0 : StringFormat("MirrorTaskManager::AddTaskInfo streamId[%u] not found after callback", taskInfo->streamId_));
70 : }
71 4 : entryPtr = &queueIt->second;
72 : }
73 10106 : auto& slot = entryPtr->queue->GetAndUpdate();
74 10106 : slot = std::move(taskInfo);
75 10106 : entryPtr->taskNum++;
76 20212 : return;
77 10106 : }
78 :
79 0 : HcclResult MirrorTaskManager::AddTaskInfo(u32 streamId, u32 taskId, u32 remoteRankId,
80 : const TaskParam &taskParam,
81 : std::shared_ptr<DfxOpInfo> dfxOpInfo, bool isMaster)
82 : {
83 0 : bool needCallback = false;
84 0 : std::unique_lock<std::mutex> lock(profMutex);
85 0 : if (dfxOpInfo == nullptr) {
86 0 : dfxOpInfo = currDfxOpInfo_;
87 : }
88 :
89 0 : auto emplaceResult = streamQueues_.emplace(streamId, MirrorStreamQueueEntry{nullptr, QueueType::Vector_Queue, 0});
90 0 : MirrorStreamQueueEntry *entryPtr = &emplaceResult.first->second;
91 0 : if (emplaceResult.second) {
92 0 : entryPtr->queueType = GetQueueType();
93 0 : entryPtr->queue = &(globalMirrorTasks_->CreateQueue(devId_, streamId, entryPtr->queueType));
94 : }
95 0 : if (UNLIKELY(entryPtr->taskNum == entryPtr->queue->Capacity())) {
96 0 : needCallback = true;
97 0 : entryPtr->taskNum = 0;
98 : }
99 :
100 0 : if (needCallback && fullyCallBack_ != nullptr) {
101 0 : lock.unlock();
102 0 : fullyCallBack_();
103 0 : lock.lock();
104 0 : auto queueIt = streamQueues_.find(streamId);
105 0 : if (queueIt == streamQueues_.end()) {
106 0 : HCCL_ERROR("[MirrorTaskManager][AddTaskInfo] streamId[%u] not found after callback", streamId);
107 0 : return HCCL_E_INTERNAL;
108 : }
109 0 : entryPtr = &queueIt->second;
110 : }
111 :
112 0 : auto& slot = entryPtr->queue->GetAndUpdate();
113 0 : if (UNLIKELY(slot == nullptr)) {
114 0 : slot = std::make_unique<TaskInfo>(streamId, taskId, remoteRankId, taskParam, dfxOpInfo, isMaster);
115 : } else {
116 0 : slot->streamId_ = streamId;
117 0 : slot->taskId_ = taskId;
118 0 : slot->taskParam_ = taskParam;
119 0 : slot->dfxOpInfo_ = dfxOpInfo;
120 0 : slot->remoteRank_ = remoteRankId;
121 0 : slot->isMaster_ = isMaster;
122 0 : slot->channelHandle_ = INVALID_U64;
123 0 : slot->getRemoteRankByHandle_ = nullptr;
124 : }
125 0 : entryPtr->taskNum++;
126 0 : return HCCL_SUCCESS;
127 0 : }
128 :
129 108 : bool MirrorTaskManager::IsStaticGraphMode(const CollOperator &collOperator) const
130 : {
131 108 : return (collOperator.staticAddr == false) && (collOperator.staticShape == false);
132 : }
133 :
134 108 : void MirrorTaskManager::SetCurrDfxOpInfo(std::shared_ptr<DfxOpInfo> dfxOpInfo)
135 : {
136 108 : if (dfxOpInfo == nullptr) {
137 0 : HCCL_ERROR("[MirrorTaskManager][SetCurrDfxOpInfo]fail, dfxOpInfo is nullptr");
138 0 : return;
139 : }
140 108 : isStaticGraphMode_ = IsStaticGraphMode(dfxOpInfo->op_);
141 108 : opMode_ = dfxOpInfo->op_.opMode;
142 108 : currDfxOpInfo_ = std::move(dfxOpInfo);
143 324 : HCCL_INFO("[MirrorTaskManager][SetCurrDfxOpInfo] Succeed, currDfxOpInfo_[%p], this[%p] !", currDfxOpInfo_.get(), this);
144 108 : return;
145 : }
146 :
147 59 : std::shared_ptr<DfxOpInfo> MirrorTaskManager::GetCurrDfxOpInfo() const
148 : {
149 59 : return currDfxOpInfo_;
150 : }
151 :
152 3 : TaskInfoQueue *MirrorTaskManager::GetQueue(u32 streamId) const
153 : {
154 3 : auto it = streamQueues_.find(streamId);
155 3 : if (it == streamQueues_.end()) {
156 2 : THROW<InternalException>(StringFormat("MirrorTaskManager::GetQueue streamId(sqId)[%u] out of range", streamId));
157 : }
158 4 : return it->second.queue;
159 : }
160 :
161 17 : std::unordered_map<u32, MirrorStreamQueueEntry>::iterator MirrorTaskManager::Begin()
162 : {
163 17 : return streamQueues_.begin();
164 : }
165 :
166 27 : std::unordered_map<u32, MirrorStreamQueueEntry>::iterator MirrorTaskManager::End()
167 : {
168 27 : return streamQueues_.end();
169 : }
170 :
171 389 : MirrorTaskManager::~MirrorTaskManager()
172 : {
173 389 : }
174 :
175 : } // namespace Hccl
|
