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