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 <cmath>
12 : #include "broadcast_oneshot_pub.h"
13 : #include "alg_template_register.h"
14 :
15 : namespace hccl {
16 0 : BroadcastHD::BroadcastHD(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
17 :
18 0 : BroadcastHD::~BroadcastHD() {}
19 :
20 0 : HcclResult BroadcastHD::Prepare(
21 : DeviceMem& inputMem, DeviceMem& outputMem, DeviceMem& scratchMem, const u64 count, const HcclDataType dataType,
22 : const Stream& stream, const HcclReduceOp reductionOp, const u32 root, std::vector<Stream>& meshStreams,
23 : const std::vector<std::shared_ptr<LocalNotify>>& meshSignal,
24 : const std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux, u32 interRank, const HcomCollOpInfo* opInfo)
25 : {
26 0 : localRank_ = interRank;
27 0 : meshStreams_ = meshStreams;
28 0 : meshSignalPtr_ = &meshSignal;
29 0 : meshSignalAuxPtr_ = &meshSignalAux;
30 0 : opInfo_ = opInfo;
31 0 : return AlgTemplateBase::Prepare(inputMem, outputMem, scratchMem, count, dataType, stream, reductionOp, root);
32 : }
33 :
34 0 : HcclResult BroadcastHD::MainRecordSub()
35 : {
36 0 : for (u32 signalIndex = 0; signalIndex < meshSignalAuxPtr_->size(); signalIndex++) {
37 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAuxPtr_)[signalIndex], profilerInput_.stage));
38 : }
39 0 : return HCCL_SUCCESS;
40 : }
41 :
42 0 : HcclResult BroadcastHD::SubWaitMain()
43 : {
44 0 : for (u32 streamIndex = 0; streamIndex < meshSignalAuxPtr_->size(); streamIndex++) {
45 0 : CHK_RET(LocalNotify::Wait(
46 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
47 : }
48 0 : return HCCL_SUCCESS;
49 : }
50 :
51 0 : HcclResult BroadcastHD::MainWaitSub()
52 : {
53 0 : for (u32 signalIndex = 0; signalIndex < meshSignalPtr_->size(); signalIndex++) {
54 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[signalIndex], profilerInput_.stage));
55 : }
56 0 : return HCCL_SUCCESS;
57 : }
58 :
59 0 : HcclResult BroadcastHD::SubRecordMain()
60 : {
61 0 : for (u32 streamIndex = 0; streamIndex < meshSignalPtr_->size(); streamIndex++) {
62 0 : CHK_RET(LocalNotify::Post(
63 : meshStreams_[streamIndex], dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
64 : }
65 0 : return HCCL_SUCCESS;
66 : }
67 :
68 0 : HcclResult BroadcastHD::PrepareStep(u32 rankSize)
69 : {
70 : u32 step;
71 0 : for (u32 rank = 0; rank < rankSize; rank++) {
72 0 : step = (rank == root_) ? 0 : static_cast<u32>(log2((rank - root_ + rankSize) % rankSize));
73 0 : stepMap_[rank] = step;
74 : }
75 :
76 0 : return HCCL_SUCCESS;
77 : }
78 :
79 : // 算法的函数入口
80 0 : HcclResult BroadcastHD::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
81 : {
82 0 : HcclResult ret = HCCL_SUCCESS;
83 0 : CHK_SMART_PTR_NULL(dispatcher_);
84 0 : CHK_PTR_NULL(stream_.ptr());
85 0 : HCCL_INFO(
86 : "BroadcastHD run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
87 : inputMem_.ptr(), outputMem_.ptr(), count_);
88 :
89 0 : if (links.size() < rankSize) {
90 0 : HCCL_ERROR(
91 : "[BroadcastHD][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(), rankSize);
92 0 : return HCCL_E_INTERNAL;
93 : }
94 :
95 0 : if (meshStreams_.size() < 1) {
96 0 : HCCL_ERROR(
97 : "[BroadcastHD][RunAsync]rank[%u] meshStreams_[%llu] is less than need[1]", rank, meshStreams_.size());
98 0 : return HCCL_E_INTERNAL;
99 : }
100 :
101 0 : CHK_RET(PrepareStep(rankSize));
102 :
103 0 : emptyMem_ = outputMem_.range(0, 0);
104 0 : nSteps_ = static_cast<u32>(log2(rankSize * base - 1));
105 :
106 0 : for (u32 step = stepMap_[rank]; step < nSteps_ - 1; step++) {
107 0 : if (step == stepMap_[rank]) {
108 0 : if (step != 0) {
109 0 : ret = RunReceive(rank, step, rankSize, links);
110 0 : CHK_PRT_RET(
111 : ret != HCCL_SUCCESS,
112 : HCCL_ERROR(
113 : "[BroadcastHD][RunAsync]rank[%u] count[%llu] step [%llu] failed in RunReceive step", rank,
114 : count_, step),
115 : ret);
116 0 : } else if (rank != root_) {
117 0 : ret = RunReceiveFirst(rank, rankSize, links);
118 0 : CHK_PRT_RET(
119 : ret != HCCL_SUCCESS,
120 : HCCL_ERROR(
121 : "[BroadcastHD][RunAsync]rank[%u] count[%llu] step [%llu] failed in RunReceiveFirst step", rank,
122 : count_, step),
123 : ret);
124 : } else {
125 0 : ret = RunSendFirst(rank, rankSize, links);
126 0 : CHK_PRT_RET(
127 : ret != HCCL_SUCCESS,
128 : HCCL_ERROR(
129 : "[BroadcastHD][RunAsync]rank[%u] count[%llu] step [%llu] failed in RunSendFirst step", rank,
130 : count_, step),
131 : ret);
132 : }
133 : } else {
134 0 : ret = RunSend(rank, step, rankSize, links);
135 0 : CHK_PRT_RET(
136 : ret != HCCL_SUCCESS,
137 : HCCL_ERROR(
138 : "[BroadcastHD][RunAsync]rank[%u] count[%llu] step [%llu] failed in RunSend step", rank, count_,
139 : step),
140 : ret);
141 : }
142 : }
143 0 : ret = RunFinalStep(rank, rankSize, links);
144 0 : CHK_PRT_RET(
145 : ret != HCCL_SUCCESS,
146 : HCCL_ERROR("[BroadcastHD][RunAsync]rank[%u] count[%llu]failed in RunFinalStep", rank, count_), ret);
147 0 : HCCL_INFO("BroadcastHD finished: rank[%u] ranksize[%u].", rank, rankSize);
148 0 : return HCCL_SUCCESS;
149 : }
150 :
151 0 : HcclResult BroadcastHD::RunFinalStep(u32 rank, u32 rankSize, const std::vector<LINK>& links)
152 : {
153 0 : HcclResult ret = HCCL_SUCCESS;
154 0 : u32 half = static_cast<u32>(pow(2, nSteps_ - 1));
155 0 : u32 logicRank = (rank - root_ + rankSize) % rankSize;
156 0 : if ((logicRank % half) < (rankSize - half)) {
157 0 : if (stepMap_[rank] == (nSteps_ - 1)) {
158 0 : ret = RunReceive(rank, nSteps_ - 1, rankSize, links);
159 0 : CHK_PRT_RET(
160 : ret != HCCL_SUCCESS,
161 : HCCL_ERROR(
162 : "[BroadcastHD][RunAsync]rank[%u] count[%llu] step [%llu] failed in RunReceive step", rank, count_,
163 : nSteps_ - 1),
164 : ret);
165 : } else {
166 0 : ret = RunSend(rank, nSteps_ - 1, rankSize, links);
167 0 : CHK_PRT_RET(
168 : ret != HCCL_SUCCESS,
169 : HCCL_ERROR(
170 : "[BroadcastHD][RunAsync]rank[%u] count[%llu] step [%llu] failed in RunSend step", rank, count_,
171 : nSteps_ - 1),
172 : ret);
173 : }
174 : } else {
175 0 : u32 unitSize = SIZE_TABLE[dataType_];
176 0 : DeviceMem userMemIn = DeviceMem::create(opInfo_->inputAddr, count_ * unitSize);
177 0 : DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), count_ * unitSize);
178 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, userMemIn, commMemOut, stream_));
179 0 : HCCL_INFO("final local cpy step %llu, rank %llu", nSteps_ - 1, rank);
180 0 : }
181 0 : return HCCL_SUCCESS;
182 : }
183 :
184 0 : u32 BroadcastHD::GetDstRank(u32 rank, u32 step, u32 rankSize)
185 : {
186 0 : u32 logicRank = (rank - root_ + rankSize) % rankSize;
187 0 : u32 logicDstRank = logicRank ^ (1 << step);
188 0 : return (logicDstRank + root_) % rankSize;
189 : }
190 :
191 0 : HcclResult BroadcastHD::RunSend(u32 rank, u32 step, u32 rankSize, const std::vector<LINK>& links)
192 : {
193 0 : u32 dstRank = GetDstRank(rank, step, rankSize);
194 0 : HCCL_INFO("RunSend: rank[%u] dstRank[%u] step [%u] count[%llu].", rank, dstRank, step, count_);
195 : // 数据准备
196 0 : u32 unitSize = SIZE_TABLE[dataType_];
197 :
198 0 : if (step == (nSteps_ - 1)) {
199 0 : CHK_RET(MainRecordSub());
200 0 : CHK_RET(SubWaitMain());
201 0 : if (rank != root_) {
202 0 : DeviceMem userMemIn = DeviceMem::create(opInfo_->inputAddr, count_ * unitSize);
203 0 : DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), count_ * unitSize);
204 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, userMemIn, commMemOut, meshStreams_[0]));
205 0 : }
206 : }
207 :
208 0 : CHK_RET(links[dstRank]->TxAck(stream_));
209 0 : CHK_RET(links[dstRank]->RxDataSignal(stream_));
210 :
211 0 : if (step == (nSteps_ - 1)) {
212 0 : CHK_RET(SubRecordMain());
213 0 : CHK_RET(MainWaitSub());
214 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem_, emptyMem_, stream_));
215 : }
216 0 : return HCCL_SUCCESS;
217 : }
218 :
219 0 : HcclResult BroadcastHD::RunReceive(u32 rank, u32 step, u32 rankSize, const std::vector<LINK>& links)
220 : {
221 0 : u32 dstRank = GetDstRank(rank, step, rankSize);
222 0 : HCCL_INFO("RunReceive: rank[%u] step[%u] outputMem[%p] count[%llu].", rank, step, outputMem_.ptr(), count_);
223 :
224 : // 数据准备
225 0 : u32 unitSize = SIZE_TABLE[dataType_];
226 0 : DeviceMem dst;
227 0 : if (step == nSteps_ - 1) {
228 0 : dst = DeviceMem::create(opInfo_->inputAddr, count_ * unitSize);
229 : } else {
230 0 : dst = outputMem_.range(0, count_ * unitSize);
231 : }
232 :
233 0 : CHK_RET(links[dstRank]->RxAck(stream_));
234 0 : void* remMemPtr = nullptr;
235 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
236 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(remMemPtr), count_ * unitSize);
237 0 : CHK_RET(HcclD2DMemcpyAsync(
238 : dispatcher_, dst, src, stream_, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
239 0 : CHK_RET(links[dstRank]->TxDataSignal(stream_));
240 0 : return HCCL_SUCCESS;
241 0 : }
242 :
243 0 : HcclResult BroadcastHD::RunSendFirst(u32 rank, u32 rankSize, const std::vector<LINK>& links)
244 : {
245 0 : u32 dstRank = GetDstRank(rank, 0, rankSize);
246 0 : HCCL_INFO("RunSendFirst: rank[%u] dstRank[%u] count[%llu].", rank, dstRank, count_);
247 : // 数据准备
248 0 : u32 unitSize = SIZE_TABLE[dataType_];
249 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, emptyMem_, emptyMem_, stream_));
250 0 : CHK_RET(MainRecordSub());
251 0 : CHK_RET(SubWaitMain());
252 :
253 0 : DeviceMem userMemIn = DeviceMem::create(opInfo_->inputAddr, count_ * unitSize);
254 0 : DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), count_ * unitSize);
255 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, commMemOut, userMemIn, meshStreams_[0]));
256 :
257 0 : CHK_RET(links[dstRank]->RxAck(stream_));
258 0 : void* remMemPtr = nullptr;
259 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
260 0 : DeviceMem dst = DeviceMem::create(static_cast<u8*>(remMemPtr), count_ * unitSize);
261 0 : CHK_RET(HcclD2DMemcpyAsync(
262 : dispatcher_, dst, userMemIn, stream_, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
263 0 : CHK_RET(links[dstRank]->TxDataSignal(stream_));
264 :
265 0 : CHK_RET(SubRecordMain());
266 0 : CHK_RET(MainWaitSub());
267 :
268 0 : return HCCL_SUCCESS;
269 0 : }
270 :
271 0 : HcclResult BroadcastHD::RunReceiveFirst(u32 rank, u32 rankSize, const std::vector<LINK>& links)
272 : {
273 0 : u32 dstRank = GetDstRank(rank, 0, rankSize);
274 0 : HCCL_INFO("RunReceiveFirst: rank[%u] dstRank[%u] count[%llu].", rank, dstRank, count_);
275 : // 数据准备
276 0 : u32 unitSize = SIZE_TABLE[dataType_];
277 :
278 0 : CHK_RET(links[dstRank]->TxAck(stream_));
279 0 : CHK_RET(links[dstRank]->RxDataSignal(stream_));
280 :
281 0 : if (nSteps_ == 1) {
282 0 : DeviceMem userMemIn = DeviceMem::create(opInfo_->inputAddr, count_ * unitSize);
283 0 : DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), count_ * unitSize);
284 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, userMemIn, commMemOut, stream_));
285 0 : }
286 0 : return HCCL_SUCCESS;
287 : }
288 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_BROADCAST_HD, BroadcastHD);
289 : } // namespace hccl
|
