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