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